diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..fe9e1ab --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +patches/**/*.patch -whitespace diff --git a/.gitignore b/.gitignore index ca3edc6..c220675 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,9 @@ dist/ checkpoint/ checkpoints/ *.safetensors +eval/results/ +/goldens/ +/Testing/ # macOS resource forks ._* diff --git a/CMakeLists.txt b/CMakeLists.txt index 28174de..7c6c01b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,9 +1,12 @@ cmake_minimum_required(VERSION 3.16) project(robotcpp VERSION 0.1.0 LANGUAGES C CXX) +include(CTest) + option(ROBOT_CPP_BUILD_ROBOT_SERVER "Build model-server target" ON) option(ROBOT_CPP_BUILD_MODEL_CLI "Build model-cli target" OFF) option(ROBOT_CPP_BUILD_ROBOT_CLIENT "Build C++ robot client targets" OFF) +option(ROBOT_CPP_BUILD_STARVLA "Build the StarVLA runtime (requires llama.cpp overlay)" OFF) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) @@ -12,6 +15,27 @@ set(CMAKE_CXX_EXTENSIONS OFF) if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/CMakeLists.txt") message(FATAL_ERROR "third_party/llama.cpp is required; run `git submodule update --init --recursive`") endif() +if(ROBOT_CPP_BUILD_STARVLA) + file(READ + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/tools/mtmd/models/qwen3vl.cpp" + ROBOT_CPP_QWEN3VL_MTMD_SOURCE) + file(READ + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/include/llama.h" + ROBOT_CPP_LLAMA_PUBLIC_HEADER) + string(FIND "${ROBOT_CPP_QWEN3VL_MTMD_SOURCE}" + "FFN_GELU_ERF" ROBOT_CPP_QWEN3VL_PARITY_PATCH_INDEX) + string(FIND "${ROBOT_CPP_LLAMA_PUBLIC_HEADER}" + "llama_set_backend_native_graphs_enabled" ROBOT_CPP_LLAMA_GRAPH_PATCH_INDEX) + if(ROBOT_CPP_QWEN3VL_PARITY_PATCH_INDEX EQUAL -1 OR + ROBOT_CPP_LLAMA_GRAPH_PATCH_INDEX EQUAL -1) + message(FATAL_ERROR + "StarVLA requires the pinned llama.cpp overlay. Run " + "`./tools/apply_patches.sh` from the repository root, " + "then configure again.") + endif() + unset(ROBOT_CPP_QWEN3VL_MTMD_SOURCE) + unset(ROBOT_CPP_LLAMA_PUBLIC_HEADER) +endif() set(LLAMA_BUILD_COMMON ON CACHE BOOL "" FORCE) set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "" FORCE) set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) @@ -20,7 +44,16 @@ add_subdirectory(third_party/llama.cpp EXCLUDE_FROM_ALL) if(NOT TARGET ggml OR NOT TARGET llama) message(FATAL_ERROR "llama.cpp must provide ggml and llama targets") endif() - +if(ROBOT_CPP_BUILD_STARVLA) + # mtmd normally inherits this variable when llama.cpp builds all tools. + if(NOT LLAMA_INSTALL_VERSION) + set(LLAMA_INSTALL_VERSION ${PROJECT_VERSION}) + endif() + add_subdirectory(third_party/llama.cpp/tools/mtmd EXCLUDE_FROM_ALL) + if(NOT TARGET mtmd) + message(FATAL_ERROR "llama.cpp must provide the mtmd target for Qwen-VL") + endif() +endif() set(ROBOT_CPP_LLAMA_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/include @@ -30,14 +63,29 @@ set(ROBOT_CPP_LLAMA_INCLUDE_DIRS ) set(SMOLVLA_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src/models/smolvla) +set(STARVLA_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src/models/starvla) set(ROBOT_SERVER_DIR ${CMAKE_CURRENT_SOURCE_DIR}/robot_server) set(ROBOT_CLIENT_DIR ${CMAKE_CURRENT_SOURCE_DIR}/robot_client) -add_library(smolvla_runtime STATIC +add_library(robotcpp_model_common STATIC src/models/ggml_backend.cpp src/models/ggml_backend.h src/models/gguf_loader.cpp src/models/gguf_loader.h + src/models/model_type.cpp +) +target_include_directories(robotcpp_model_common + PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} +) +target_link_libraries(robotcpp_model_common PUBLIC ggml) +target_compile_features(robotcpp_model_common PUBLIC cxx_std_17) +if(NOT MSVC) + target_compile_options(robotcpp_model_common PRIVATE -Wno-cast-qual) +endif() + +add_library(smolvla_runtime STATIC ${SMOLVLA_DIR}/smolvla_engine.cpp ${SMOLVLA_DIR}/smolvla_engine.h ${SMOLVLA_DIR}/state_proj.cpp @@ -53,17 +101,13 @@ target_include_directories(smolvla_runtime ${SMOLVLA_DIR} ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} ) -target_link_libraries(smolvla_runtime PUBLIC ggml llama) +target_link_libraries(smolvla_runtime PUBLIC robotcpp_model_common llama) target_compile_features(smolvla_runtime PUBLIC cxx_std_17) if(NOT MSVC) target_compile_options(smolvla_runtime PRIVATE -Wno-cast-qual) endif() add_library(pi0_engine STATIC - src/models/ggml_backend.cpp - src/models/ggml_backend.h - src/models/gguf_loader.cpp - src/models/gguf_loader.h src/models/pi0/types.h src/models/pi0/action.cpp src/models/pi0/action.h @@ -87,13 +131,65 @@ target_include_directories(pi0_engine ${CMAKE_CURRENT_SOURCE_DIR}/src ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} ) -target_link_libraries(pi0_engine PUBLIC ggml llama) +target_link_libraries(pi0_engine PUBLIC robotcpp_model_common llama) target_compile_features(pi0_engine PUBLIC cxx_std_17) if(NOT MSVC) target_compile_options(pi0_engine PRIVATE -Wno-cast-qual) endif() +if(ROBOT_CPP_BUILD_STARVLA) +add_library(starvla_runtime STATIC + ${STARVLA_DIR}/fast_codec.cpp + ${STARVLA_DIR}/fast_codec.h + ${STARVLA_DIR}/fast_policy.cpp + ${STARVLA_DIR}/fast_policy.h + ${STARVLA_DIR}/groot_policy.cpp + ${STARVLA_DIR}/groot_policy.h + ${STARVLA_DIR}/groot_prompt.cpp + ${STARVLA_DIR}/groot_prompt.h + ${STARVLA_DIR}/normalization.cpp + ${STARVLA_DIR}/normalization.h + ${STARVLA_DIR}/oft_image_preprocess.cpp + ${STARVLA_DIR}/oft_image_preprocess.h + ${STARVLA_DIR}/oft_prompt.cpp + ${STARVLA_DIR}/oft_prompt.h + ${STARVLA_DIR}/oft_policy.cpp + ${STARVLA_DIR}/oft_policy.h + ${STARVLA_DIR}/pi_policy.cpp + ${STARVLA_DIR}/pi_policy.h + ${STARVLA_DIR}/pi_v3_policy.cpp + ${STARVLA_DIR}/pi_v3_policy.h + ${STARVLA_DIR}/policy_gguf.h + ${STARVLA_DIR}/qwen3vl_bridge.cpp + ${STARVLA_DIR}/qwen3vl_bridge.h + ${STARVLA_DIR}/starvla_engine.cpp + ${STARVLA_DIR}/starvla_engine.h + third_party/llama.cpp/examples/gguf-hash/deps/sha256/sha256.c +) +target_include_directories(starvla_runtime + PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${ROBOT_CPP_LLAMA_INCLUDE_DIRS} + PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/vendor + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/examples/gguf-hash/deps + ${CMAKE_CURRENT_SOURCE_DIR}/third_party/llama.cpp/examples/gguf-hash/deps/sha256 +) +target_link_libraries(starvla_runtime PUBLIC robotcpp_model_common llama mtmd) +target_compile_features(starvla_runtime PUBLIC cxx_std_17) +if(GGML_CUDA) + enable_language(CUDA) + target_sources(starvla_runtime PRIVATE ${STARVLA_DIR}/qwen_bf16_round_cuda.cu) + target_compile_definitions(starvla_runtime PRIVATE ROBOTCPP_STARVLA_CUDA=1) + set_property(TARGET starvla_runtime PROPERTY CUDA_STANDARD 17) +endif() +if(NOT MSVC) + target_compile_options(starvla_runtime PRIVATE -Wno-cast-qual) +endif() +endif() + add_library(robotcpp STATIC + src/models/argument_parse.h src/models/model.h src/models/model_factory.cpp src/models/pi0/pi0_model.cpp @@ -108,6 +204,19 @@ target_include_directories(robotcpp ) target_link_libraries(robotcpp PUBLIC smolvla_runtime pi0_engine) target_compile_features(robotcpp PUBLIC cxx_std_17) +if(ROBOT_CPP_BUILD_STARVLA) + target_sources(robotcpp PRIVATE + ${STARVLA_DIR}/starvla_model.cpp + ${STARVLA_DIR}/starvla_model.h) + target_link_libraries(robotcpp PUBLIC starvla_runtime) + target_compile_definitions(robotcpp PUBLIC ROBOT_CPP_BUILD_STARVLA=1) +endif() + +if(BUILD_TESTING AND ROBOT_CPP_BUILD_STARVLA) + add_executable(robotcpp-starvla-model-test tests/starvla/model_test.cpp) + target_link_libraries(robotcpp-starvla-model-test PRIVATE robotcpp) + add_test(NAME robotcpp-starvla-model-test COMMAND robotcpp-starvla-model-test) +endif() if(ROBOT_CPP_BUILD_ROBOT_SERVER OR ROBOT_CPP_BUILD_ROBOT_CLIENT) add_library(robot_server_common STATIC @@ -139,6 +248,7 @@ if(ROBOT_CPP_BUILD_ROBOT_CLIENT) set_target_properties(model-cpp-client-example PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) target_link_libraries(model-cpp-client-example PRIVATE model_client_cpp) target_compile_features(model-cpp-client-example PRIVATE cxx_std_17) + endif() if(ROBOT_CPP_BUILD_ROBOT_SERVER) @@ -165,7 +275,6 @@ if(ROBOT_CPP_BUILD_ROBOT_SERVER) ) target_link_libraries(model-server PRIVATE robot_server_core robotcpp) target_compile_features(model-server PRIVATE cxx_std_17) - add_executable(smolvla-raw-predict ${ROBOT_SERVER_DIR}/test/smolvla_raw_predict.cpp) set_target_properties(smolvla-raw-predict PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin) target_include_directories(smolvla-raw-predict PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src ${SMOLVLA_DIR}) diff --git a/README.md b/README.md index 055e826..4e02885 100644 --- a/README.md +++ b/README.md @@ -54,8 +54,13 @@ We also provide two tools to support robot model development: git clone https://github.com/Robot-cpp/robot.cpp cd robot.cpp git submodule update --init --recursive +./tools/apply_patches.sh ``` +The launch scripts below configure and build `model-server` automatically. For +a manual StarVLA build, enable `ROBOT_CPP_BUILD_STARVLA`; see the +[Robot Server build instructions](robot_server/README.md#manual-build). + This section introduces three usage paths to help you quickly understand the repository: * Starting `model-server` and connecting it to a minimal dummy `model-client`. @@ -95,6 +100,14 @@ After downloading, run `model-server` like this: For general local setups, we provide ready-to-use build-and-launch shells for three platforms. You can modify the environment variables inside the scripts, or override them directly with `export`. See [robot_server/README.md](robot_server/README.md) for details. +For example, from the repository root on Linux with CUDA: + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/smolvla-so101-fp32 +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + | Backend | macOS | Linux | Windows | | ------- | ------------------------------------------------------- | -------------------------------------------------------- | ----------------------------------------------------------- | | CUDA | - | `robot_server/shell/launch_robot_server_linux_cuda.sh` | `robot_server/shell/launch_robot_server_windows_cuda.bat` | @@ -127,7 +140,7 @@ We provide a build-to-run example in `robot_client/shell/cpp_client_example.sh`. | `ROBOT_CPP_ROOT` | unset; required | Repository root. | | `BUILD_DIR` | `${ROBOT_CPP_ROOT}/build_robot_client` | C++ client CMake build directory. | | `PORT` | `5555` | Server port used by the client. | -| `BUILD_CLIENT` | `0` | Whether to force rebuild the client. Set to`1` to rebuild even if the binary already exists. | +| `BUILD_CLIENT` | `0` | Whether to force rebuild the client. Set to `1` to rebuild even if the binary already exists. | | `CMAKE_BIN` | `cmake` | CMake command path, useful for selecting a custom CMake binary. | Then run: @@ -138,7 +151,9 @@ bash robot_client/shell/cpp_client_example.sh ### 🧪 Using model-server in simulation, using LIBERO as the example -See the [LIBERO simulation evaluation guide](eval/libero/README.md). +See the [LIBERO simulation evaluation guide](eval/libero/README.md). To run +local StarVLA GGUF models on the WidowX Bridge tasks, see the +[SimplerEnv Bridge guide](eval/simpler_env/README.md). ### 🦾 Using model-server on real hardware, using SO-101 as the example @@ -150,20 +165,33 @@ See the [SO-101 deployment guide](eval/lerobot_so101/README.md). ## ⚡ Performance -We benchmark Robot.cpp on several platforms. Each measurement uses 5 warmup runs and 100 loop runs. The reported latency is the average time from receiving the image, through preprocessing and forward inference, to producing a usable action chunk, measured in milliseconds. All state projectors remain in f32 precision. +We benchmark Robot.cpp on several platforms. Each measurement uses 5 warmup runs and 100 loop runs. The reported latency is the average time from receiving the image, through preprocessing and forward inference, to producing a usable action chunk, measured in milliseconds. State projectors, where present, remain in f32 precision. For the LIBERO setting, the input contains two 256x256 images and an 8-dimensional state. For the SO-101 real-robot setting, the input contains one 224x224 image and a 6-dimensional state. For SmolVLA preprocessing, we follow the official default setting: images are first resized to 512x512. -| Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | -| ---------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | -| smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | -| smolvla@libero (f32) | 577 | 236 | 32 | 142 | 42 | 299 | -| smolvla@so-101 (bf16*) | 339 | 145 | 23 | 77 | 36 | 184 | -| smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | -| pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | -| pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +For StarVLA, the input contains one 224x224 image and no robot state. Qwen and +the multimodal projector use bf16; OFT, GR00T, PI, and PI_v3 policies use f32. +FAST stores its action codec in the policy GGUF. +The StarVLA A100 results use an A100-PCIE-40GB with 8 CPU threads, +`n_ctx=2048`, `n_batch=2048`, and noise seed 0. + +| Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | +| ----------------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | +| smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | +| smolvla@libero (f32) | 577 | 236 | 32 | 142 | 42 | 299 | +| smolvla@so-101 (bf16*) | 339 | 145 | 23 | 77 | 36 | 184 | +| smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | +| pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | +| pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +| starvla/oft@bridge | - | - | - | - | 50 | - | +| starvla/groot@bridge | - | - | - | - | 54 | - | +| starvla/pi_v3@bridge | - | - | - | - | 112 | - | +| starvla/qwen25_oft@bridge | - | - | - | - | 42 | - | +| starvla/qwen25_groot@bridge | - | - | - | - | 51 | - | +| starvla/qwen25_pi@bridge | - | - | - | - | 101 | - | +| starvla/qwen25_fast@bridge | - | - | - | - | 386 | - | > `bf16*`: on Mac, f16 results are used in place of bf16 because current Mac bf16 support is not ideal. > `OOM/offload`: pi0@libero (f32) runs out of memory on RTX 3060 and triggers offload, so we do not report a latency number for now. @@ -230,6 +258,55 @@ This section lists converted GGUF models that can be used directly with `model-s f32 pi0-libero-f32 + + StarVLA Qwen3-VL OFT + Bridge + StarVLA/Qwen3VL-OFT-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen3-oft-bridge-bf16 + + + StarVLA Qwen3-VL GR00T + Bridge + StarVLA/Qwen3VL-GR00T-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen3-groot-bridge-bf16 + + + StarVLA Qwen3-VL PI_v3 + Bridge + StarVLA/Qwen3VL-PI_v3-Bridge-RT_1 + bf16 + f32 policy + starvla-qwen3-pi-v3-bridge-bf16 + + + StarVLA Qwen2.5-VL OFT + Bridge + StarVLA/Qwen-OFT-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-oft-bridge-bf16 + + + StarVLA Qwen2.5-VL GR00T + Bridge + StarVLA/Qwen-GR00T-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-groot-bridge-bf16 + + + StarVLA Qwen2.5-VL PI + Bridge + StarVLA/Qwen-PI-Bridge-RT-1 + bf16 + f32 policy + starvla-qwen25-pi-bridge-bf16 + + + StarVLA Qwen2.5-VL FAST + Bridge + StarVLA/Qwen-FAST-Bridge-RT-1 + bf16 + codec + starvla-qwen25-fast-bridge-bf16 + @@ -269,6 +346,7 @@ robot.cpp/ ├── eval/ │ ├── base_platform.py # Shared base class for real-robot platforms │ ├── libero/ # LIBERO simulation evaluation +│ ├── simpler_env/ # SimplerEnv WidowX / Bridge evaluation │ └── lerobot_so101/ # SO-101 real-robot scripts and examples └── third_party/ ├── llama.cpp/ # ggml / llama.cpp backend @@ -325,6 +403,7 @@ Robot.cpp's design and implementation benefit from several excellent open-source * [llama.cpp](https://github.com/ggerganov/llama.cpp): provides lightweight local inference, the GGML/GGUF ecosystem, and cross-platform backend foundations. This project continues building robot model inference capabilities on top of its engineering philosophy and low-level runtime. * [LeRobot](https://github.com/huggingface/lerobot): provides reference implementations for robot data, policy training, and real-robot integration. The SO-101 real-robot example and parts of the evaluation flow in this project are inspired by the LeRobot ecosystem. * [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO): provides robot simulation tasks and evaluation benchmarks. The LIBERO simulation evaluation flow in this project is based on its task environments and benchmark design. +* [SimplerEnv](https://github.com/simpler-env/SimplerEnv): provides real-to-sim robot evaluation environments. The StarVLA Bridge success-rate evaluation uses its WidowX task suite and official visual-matching assets. * [OpenPI](https://github.com/Physical-Intelligence/openpi): provides the pi0 policy model and related open-source implementation. The pi0 runtime, conversion, and evaluation work in this project references OpenPI's model design. Thanks to these projects and communities for their contributions to robot learning and on-device inference. diff --git a/README_ZH.md b/README_ZH.md index 78c5e4a..0d61dfa 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -54,8 +54,13 @@ Robot.cpp是一个轻量化的on-device机器人模型推理框架，在llama.cp git clone https://github.com/Robot-cpp/robot.cpp cd robot.cpp git submodule update --init --recursive +./tools/apply_patches.sh ``` +下文的启动脚本会自动配置并编译 `model-server`。手动构建 StarVLA 时需要开启 +`ROBOT_CPP_BUILD_STARVLA`，详见 +[Robot Server 构建说明](robot_server/README_ZH.md#手动构建)。 + 我们介绍三类使用案例来帮助你快速了解本仓库： * model-server的启动，其与最小dummy model-client通信的案例。 @@ -95,6 +100,14 @@ git submodule update --init --recursive 对于更加一般的情况，我们也提供了三个平台的开箱即用编译+启动的shell，可以通过修改shell里的环境变量，或者直接export的形式来快速在本机实现启动。详情参见 [robot_server/README_ZH.md](robot_server/README_ZH.md) +例如，在 Linux CUDA 环境中从仓库根目录运行： + +```bash +export ROBOT_CPP_ROOT="$PWD" +export GGUF_DIR=/path/to/smolvla-so101-fp32 +bash robot_server/shell/launch_robot_server_linux_cuda.sh +``` + | Backend | macOS | Linux | Windows | | ------- | ------------------------------------------------------- | -------------------------------------------------------- | ----------------------------------------------------------- | | CUDA | - | `robot_server/shell/launch_robot_server_linux_cuda.sh` | `robot_server/shell/launch_robot_server_windows_cuda.bat` | @@ -127,7 +140,7 @@ python robot_client/examples/python/minimal_example.py | `ROBOT_CPP_ROOT` | 无，必须设置 | 仓库根目录。 | | `BUILD_DIR` | `${ROBOT_CPP_ROOT}/build_robot_client` | C++ client 的 CMake build 目录 | | `PORT` | `5555` | client 连接的 server port | -| `BUILD_CLIENT` | `0` | 是否强制重新build client。设为`1` 时即使 binary 已存在也会重新 build | +| `BUILD_CLIENT` | `0` | 是否强制重新build client。设为 `1` 时即使 binary 已存在也会重新 build | | `CMAKE_BIN` | `cmake` | 使用的 CMake 命令路径，可用于指定自定义 CMake | 然后运行下面的bash： @@ -138,7 +151,9 @@ bash robot_client/shell/cpp_client_example.sh ### 🧪 model-server在仿真平台上的使用（以LIBERO为例） -详见 [LIBERO 仿真评测说明](eval/libero/README_ZH.md)。 +详见 [LIBERO 仿真评测说明](eval/libero/README_ZH.md)。在 WidowX Bridge 任务上运行 +StarVLA，并比较 Python checkpoint 与 GGUF 的方法见 +[SimplerEnv Bridge 说明](eval/simpler_env/README_ZH.md)。 ### 🦾 model-server在真机平台上的使用（以SO-101为例） @@ -156,6 +171,11 @@ bash robot_client/shell/cpp_client_example.sh 其中对于smolvla的preprocess设定，参考官方的基本设定，即首先会将图片变成512*512。 +StarVLA 使用一张 224x224 图像且不输入 robot state。Qwen 和 multimodal projector +使用 bf16，OFT、GR00T、PI 和 PI_v3 policy 使用 f32；FAST 的 policy GGUF 保存 action +codec。A100 数据在 A100-PCIE-40GB、8 个 CPU 线程、`n_ctx=2048`、`n_batch=2048` 和 +noise seed 0 下测得。 + | Model | Mac M4 Pro (CPU) | Mac M4 Pro (Metal) | RTX 4090 | RTX 3060 | A100 | Jetson AGX Orin | | ---------------------- | ---------------: | -----------------: | -------: | ----------: | ---: | --------------: | | smolvla@libero (bf16*) | 527 | 216 | 28 | 116 | 43 | 282 | @@ -164,15 +184,24 @@ bash robot_client/shell/cpp_client_example.sh | smolvla@so-101 (f32) | 396 | 158 | 24 | 92 | 34 | 200 | | pi0@libero (f32) | 1839 | 710 | 83 | OOM/offload | 71 | 956 | | pi0@libero (bf16*) | 1954 | 635 | 57 | 267 | 66 | 498 | +| starvla/oft@bridge | - | - | - | - | 50 | - | +| starvla/groot@bridge | - | - | - | - | 54 | - | +| starvla/pi_v3@bridge | - | - | - | - | 112 | - | +| starvla/qwen25_oft@bridge | - | - | - | - | 42 | - | +| starvla/qwen25_groot@bridge | - | - | - | - | 51 | - | +| starvla/qwen25_pi@bridge | - | - | - | - | 101 | - | +| starvla/qwen25_fast@bridge | - | - | - | - | 386 | - | > `bf16*`：在 Mac上使用 f16 结果替代 bf16，因为当前 Mac对 bf16 的支持不够好。 > `OOM/offload`：pi0@libero (f32) 在 RTX 3060 上会 OOM 并触发 offload，因此暂时不报告 latency 数值。 --- -## 🧩 model-zoo +## 🧩 Model Zoo -这里整理一些已经转换好的 GGUF 模型，可以直接配合 `model-server` 做smoke test，以方便quick start！但针对自己的实际场景，我们推荐使用[hf2gguf](tools/hf2gguf/README_ZH.md)来生成自己的GGUF model！并且对于不同的部分，您还可以自定义不同的精度，来实现不同部分的精度组合（事实上，不同部分的最优精度通常是不同的），我们的例子中，state proj始终保持f32精度，其他的gguf随着precision精度变化而变化，您可以自行组合，探索更好更高效的性能tradeoff！ +下表列出可直接配合 `model-server` 使用的 GGUF 模型。实际部署时，建议使用 +[`hf2gguf`](tools/hf2gguf/README_ZH.md) 转换自己的 checkpoint。各组件可以分别选择 +精度；表中示例的 state projector 固定为 f32，其余组件采用标注的精度。 @@ -230,6 +259,55 @@ bash robot_client/shell/cpp_client_example.sh + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
f32 pi0-libero-f32
StarVLA Qwen3-VL OFTBridgeStarVLA/Qwen3VL-OFT-Bridge-RT-1bf16 + f32 policystarvla-qwen3-oft-bridge-bf16
StarVLA Qwen3-VL GR00TBridgeStarVLA/Qwen3VL-GR00T-Bridge-RT-1bf16 + f32 policystarvla-qwen3-groot-bridge-bf16
StarVLA Qwen3-VL PI_v3BridgeStarVLA/Qwen3VL-PI_v3-Bridge-RT_1bf16 + f32 policystarvla-qwen3-pi-v3-bridge-bf16
StarVLA Qwen2.5-VL OFTBridgeStarVLA/Qwen-OFT-Bridge-RT-1bf16 + f32 policystarvla-qwen25-oft-bridge-bf16
StarVLA Qwen2.5-VL GR00TBridgeStarVLA/Qwen-GR00T-Bridge-RT-1bf16 + f32 policystarvla-qwen25-groot-bridge-bf16
StarVLA Qwen2.5-VL PIBridgeStarVLA/Qwen-PI-Bridge-RT-1bf16 + f32 policystarvla-qwen25-pi-bridge-bf16
StarVLA Qwen2.5-VL FASTBridgeStarVLA/Qwen-FAST-Bridge-RT-1bf16 + codecstarvla-qwen25-fast-bridge-bf16
@@ -269,6 +347,7 @@ robot.cpp/ ├── eval/ │ ├── base_platform.py # 真机 platform 的统一基类 │ ├── libero/ # LIBERO 仿真评测 +│ ├── simpler_env/ # SimplerEnv WidowX / Bridge 仿真评测 │ └── lerobot_so101/ # SO-101 真机相关脚本与示例 └── third_party/ ├── llama.cpp/ # ggml / llama.cpp 后端 @@ -325,6 +404,7 @@ robot.cpp 的设计与实现受益于多个优秀的开源项目： * [llama.cpp](https://github.com/ggerganov/llama.cpp)：提供了轻量化本地推理、GGML/GGUF 生态与跨平台后端基础，本项目在其工程哲学和底层能力上继续构建机器人模型推理框架。 * [LeRobot](https://github.com/huggingface/lerobot)：提供了机器人数据、策略训练与真实机器人接入的参考实现，本项目的 SO-101 真机示例与部分评测流程参考了 LeRobot 生态。 * [LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO)：提供了机器人仿真任务与评测基准，本项目的 LIBERO 仿真评测流程基于其任务环境与 benchmark 设计。 +* [SimplerEnv](https://github.com/simpler-env/SimplerEnv)：提供了 real-to-sim 机器人评测环境，本项目的 StarVLA Bridge 成功率评测使用其 WidowX 任务集与官方 visual-matching 资产。 * [OpenPI](https://github.com/Physical-Intelligence/openpi)：提供了pi0策略模型与相关开源实现，本项目的 pi0 相关 runtime、转换与评测工作参考了 OpenPI 的模型设计。 感谢这些项目和社区为机器人学习与端侧推理生态做出的贡献。 diff --git a/eval/README.md b/eval/README.md index 9cf9e75..0d9e672 100644 --- a/eval/README.md +++ b/eval/README.md @@ -22,6 +22,7 @@ The repo root [README.md](../README.md) describes the three-layer layout: eval/ ├── base_platform.py # Shared base class for real-robot / sim platforms ├── libero/ # LIBERO sim benchmark (multi-camera, batch rollout) +├── simpler_env/ # SimplerEnv WidowX / Bridge closed-loop benchmark └── lerobot_so101/ # SO-101 real-robot sync closed-loop example ``` @@ -30,12 +31,16 @@ eval/ | Directory | Scenario | Notes | |---|---|---| | [`libero/`](libero/README.md) | Sim eval | LIBERO benchmark with C++ policy rollout and LeRobot baseline. [中文](libero/README_ZH.md) | +| [`simpler_env/`](simpler_env/README.md) | Sim eval | Runs StarVLA Python and GGUF on the SimplerEnv WidowX Bridge tasks. [中文](simpler_env/README_ZH.md) | | [`lerobot_so101/`](lerobot_so101/README.md) | Real robot | SO-101 follower + single-camera observe → predict → act loop. [中文](lerobot_so101/README_ZH.md) | The two examples are organized slightly differently: - **SO-101** follows the standard `BasePlatform` + `RobotPolicy` + `SyncControlLoop` path—use it as the template for new real-robot platforms. - **LIBERO** implements a dedicated observation adapter under `eval/libero/policy/` (multi-camera, state packing, sim rollout) and does **not** inherit `BasePlatform`—use it as a reference for new **sim benchmarks**. +- **SimplerEnv** follows the same dedicated runner pattern and implements the + WidowX action transform, temporal ensemble, normalization profile, and Bridge + task settings for both Python and C++ runs. ## Standard closed-loop data flow @@ -237,6 +242,7 @@ LIBERO’s [`ModelServerPolicy`](libero/policy/model_server.py) is an example of - [SO-101 real-robot guide](lerobot_so101/README.md) · [中文](lerobot_so101/README_ZH.md) - [LIBERO sim eval](libero/README.md) · [中文](libero/README_ZH.md) +- [SimplerEnv Bridge eval](simpler_env/README.md) · [中文](simpler_env/README_ZH.md) - [robot_server launch and protocol](../robot_server/README.md) - [robot_client and policy](../robot_client/README.md) - [Adding a new model runtime](../src/README.md) diff --git a/eval/README_ZH.md b/eval/README_ZH.md index eb1e781..99784b1 100644 --- a/eval/README_ZH.md +++ b/eval/README_ZH.md @@ -24,6 +24,7 @@ eval/ ├── base_platform.py # 真机 / 仿真 platform 的统一基类 ├── libero/ # LIBERO 仿真 benchmark（多相机、批量 rollout） +├── simpler_env/ # SimplerEnv WidowX / Bridge 闭环 benchmark └── lerobot_so101/ # SO-101 真机同步闭环示例 ``` @@ -33,6 +34,7 @@ eval/ | 目录 | 场景 | 说明 | | ---------------------------------------------- | ---- | ------------------------------------------------------------------------------------------ | | `[libero/](libero/README_ZH.md)` | 仿真评测 | 面向 LIBERO benchmark，含 C++ policy rollout 与 LeRobot baseline 对比。[English](libero/README.md) | +| `[simpler_env/](simpler_env/README_ZH.md)` | 仿真评测 | 在 SimplerEnv WidowX Bridge 任务上运行并比较 StarVLA Python 与 GGUF。[English](simpler_env/README.md) | | `[lerobot_so101/](lerobot_so101/README_ZH.md)` | 真机闭环 | SO-101 follower + 单相机的 observe → predict → act 同步控制。[English](lerobot_so101/README.md) | @@ -40,6 +42,7 @@ eval/ - **SO-101** 走标准 `BasePlatform` + `RobotPolicy` + `SyncControlLoop` 路径，适合作为新增真机 platform 的模板。 - **LIBERO** 在 `eval/libero/policy/` 里实现了专用的 observation 适配（多相机、state 拼接、仿真 rollout），不继承 `BasePlatform`，适合作为新增 **仿真 benchmark** 的参考。 +- **SimplerEnv** 沿用专用 runner 结构，为 Python 和 C++ 实现相同的 WidowX action 变换、时序集成、normalization profile 与 Bridge 任务设置。 ## 标准闭环数据流 @@ -244,6 +247,7 @@ LIBERO 的 `[ModelServerPolicy](libero/policy/model_server.py)` 即为自定义 - [SO-101 真机使用说明](lerobot_so101/README_ZH.md) - [LIBERO 仿真评测说明](libero/README_ZH.md) +- [SimplerEnv Bridge 仿真评测说明](simpler_env/README_ZH.md) - [robot_server 启动与协议](../robot_server/README_ZH.md) - [robot_client 与 policy](../robot_client/README.md) - [新增模型 runtime](../src/README_ZH.md) diff --git a/eval/simpler_env/README.md b/eval/simpler_env/README.md new file mode 100644 index 0000000..b22dc91 --- /dev/null +++ b/eval/simpler_env/README.md @@ -0,0 +1,97 @@ +# SimplerEnv WidowX Bridge Eval + +This directory evaluates the robot.cpp StarVLA GGUF runtime on the SimplerEnv +WidowX Bridge tasks. + +## Protocol + +- Four Bridge tasks with object episodes `0..23` +- At most 120 steps per episode at 5 Hz +- Visual-matching RGB overlay resized to 224x224 with OpenCV `INTER_AREA` +- One action chunk per step with the official seven-prediction adaptive ensemble + +A full run contains 96 rollouts. The result reports overall and per-task success +rates; subset runs use `partial` coverage. + +## Setup + +Convert a checkpoint as described in the +[StarVLA guide](../../tools/hf2gguf/starvla/README.md) and build the CUDA runtime. + +The environment uses these revisions: + +```text +SimplerEnv: 06accaca93535902d408da4855f21cece12bceb7 +ManiSkill2_real2sim: ef7a4d4fdf4b69f2c2154db5b15b9ac8dfe10682 +``` + +```bash +conda env create -f eval/simpler_env/environment.yaml +conda activate robotcpp-simpler-env + +git clone --recurse-submodules https://github.com/simpler-env/SimplerEnv \ + ckpts/simpler_env/source/SimplerEnv +git -C ckpts/simpler_env/source/SimplerEnv checkout \ + 06accaca93535902d408da4855f21cece12bceb7 +git -C ckpts/simpler_env/source/SimplerEnv submodule update --init --recursive + +pip install -e ckpts/simpler_env/source/SimplerEnv/ManiSkill2_real2sim +pip install -e ckpts/simpler_env/source/SimplerEnv +``` + +Headless simulation requires a working Vulkan ICD. Run SimplerEnv's environment +test first to confirm that SAPIEN can find a rendering device. + +## Run + +`VARIANT` accepts `oft`, `groot`, `pi_v3`, `qwen25_oft`, `qwen25_groot`, +`qwen25_pi`, and `qwen25_fast`. + +Run the full profile: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=oft \ +OUTPUT=ckpts/starvla/results/oft/bridge.json \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +Run one smoke episode: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=groot TASK_IDS=0 EPISODE_IDS=0 \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +The script reads three GGUF files from `ckpts/starvla/gguf/` and uses +`build_cuda/bin/model-server` by default. Common overrides are `GGUF_DIR`, +`SERVER_BIN`, `PYTHON`, `SIMPLER_ENV_ROOT`, `TASK_IDS`, +`EPISODE_IDS`, `REPEATS`, and `OUTPUT`. + +Each task/repeat starts a fresh model-server. Results include checkpoint +identity, rollout records, success rates, and timing summaries. + +## Latency + +Benchmark the official PyTorch checkpoint directly: + +```bash +CUDA_VISIBLE_DEVICES=0 python -m eval.simpler_env.runners.latency_starvla \ + --variant oft --compile-model +``` + +The runner selects the checkpoint, Qwen assets, and Bridge normalization from +the StarVLA catalog. It reports policy, action unnormalization, and total +latency after 5 warmup calls and 20 measured calls. Policy latency includes +StarVLA's image/text preprocessing and model forward. Omit `--compile-model` +for eager PyTorch. Compilation is lazy; the first FAST warmup can take several +minutes and is not included in the reported measurements. + +For the robot.cpp model-server path, use the common server benchmark: + +```bash +CUDA_VISIBLE_DEVICES=0 N_BATCH=2048 SKIP_BUILD=1 \ +GGUF_DIR="$PWD/ckpts/starvla/gguf/oft" \ +bash robot_server/test/test_server_latency.sh starvla linux-cuda starvla-bridge +``` diff --git a/eval/simpler_env/README_ZH.md b/eval/simpler_env/README_ZH.md new file mode 100644 index 0000000..4a049b2 --- /dev/null +++ b/eval/simpler_env/README_ZH.md @@ -0,0 +1,97 @@ +# SimplerEnv WidowX Bridge 评测 + +本目录使用 robot.cpp 的 StarVLA GGUF runtime 运行 SimplerEnv WidowX Bridge 任务。 + +## 评测设置 + +- 四个 Bridge 任务，每个任务包含 object episode `0..23` +- 每个 episode 最多 120 步，控制频率 5 Hz +- visual-matching RGB overlay 使用 OpenCV `INTER_AREA` 缩放到 224x224 +- 每步预测一个 action chunk，并对最近七次预测做自适应集成 + +完整评测包含 96 个 rollout。结果文件同时记录总体和各任务成功率；子集运行会标记为 +`partial` coverage。 + +## 安装 + +先按 [StarVLA 转换说明](../../tools/hf2gguf/starvla/README.md) 生成 GGUF，并完成 CUDA +构建。 + +SimplerEnv 使用以下 revision： + +```text +SimplerEnv: 06accaca93535902d408da4855f21cece12bceb7 +ManiSkill2_real2sim: ef7a4d4fdf4b69f2c2154db5b15b9ac8dfe10682 +``` + +```bash +conda env create -f eval/simpler_env/environment.yaml +conda activate robotcpp-simpler-env + +git clone --recurse-submodules https://github.com/simpler-env/SimplerEnv \ + ckpts/simpler_env/source/SimplerEnv +git -C ckpts/simpler_env/source/SimplerEnv checkout \ + 06accaca93535902d408da4855f21cece12bceb7 +git -C ckpts/simpler_env/source/SimplerEnv submodule update --init --recursive + +pip install -e ckpts/simpler_env/source/SimplerEnv/ManiSkill2_real2sim +pip install -e ckpts/simpler_env/source/SimplerEnv +``` + +无头运行需要可用的 Vulkan ICD。请先运行 SimplerEnv 自带的环境测试，确认 SAPIEN 能找到 +渲染设备。 + +## 运行 + +`VARIANT` 支持： + +```text +oft groot pi_v3 qwen25_oft qwen25_groot qwen25_pi qwen25_fast +``` + +完整运行： + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=oft \ +OUTPUT=ckpts/starvla/results/oft/bridge.json \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +快速检查一个 episode： + +```bash +CUDA_VISIBLE_DEVICES=0 \ +VARIANT=groot TASK_IDS=0 EPISODE_IDS=0 \ +bash eval/simpler_env/scripts/run_model_server.sh +``` + +脚本默认从 `ckpts/starvla/gguf/` 读取三个 GGUF，并使用 +`build_cuda/bin/model-server`。常用覆盖项包括 `GGUF_DIR`、`SERVER_BIN`、`PYTHON`、 +`SIMPLER_ENV_ROOT`、`TASK_IDS`、`EPISODE_IDS`、`REPEATS` 和 `OUTPUT`。 + +每个 task/repeat 会启动新的 model-server。结果包含 checkpoint 标识、rollout 明细、成功率 +和各阶段耗时。 + +## 延迟测试 + +直接测试官方 PyTorch checkpoint： + +```bash +CUDA_VISIBLE_DEVICES=0 python -m eval.simpler_env.runners.latency_starvla \ + --variant oft --compile-model +``` + +runner 会根据 StarVLA catalog 选择 checkpoint、Qwen 资源和 Bridge 归一化配置。默认先预热 +5 次，再统计 20 次推理，并分别报告 policy、action 反归一化和总耗时。policy 耗时包含 +StarVLA 的图像/文本预处理和模型 forward。去掉 `--compile-model` 即可测试 eager +PyTorch。`torch.compile` 为惰性编译；FAST 第一次预热可能需要几分钟，这部分不会计入 +最终统计。 + +robot.cpp model-server 使用统一的服务端测试脚本： + +```bash +CUDA_VISIBLE_DEVICES=0 N_BATCH=2048 SKIP_BUILD=1 \ +GGUF_DIR="$PWD/ckpts/starvla/gguf/oft" \ +bash robot_server/test/test_server_latency.sh starvla linux-cuda starvla-bridge +``` diff --git a/eval/simpler_env/__init__.py b/eval/simpler_env/__init__.py new file mode 100644 index 0000000..55ef27c --- /dev/null +++ b/eval/simpler_env/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation integration.""" diff --git a/eval/simpler_env/environment.yaml b/eval/simpler_env/environment.yaml new file mode 100644 index 0000000..1b17b4d --- /dev/null +++ b/eval/simpler_env/environment.yaml @@ -0,0 +1,19 @@ +name: robotcpp-simpler-env +channels: + - conda-forge +dependencies: + - python=3.10 + - pip + - ffmpeg + - pip: + - numpy==1.24.4 + - scipy==1.11.4 + - opencv-python==4.11.0.86 + - opencv-python-headless==4.11.0.86 + - setuptools<81 + - transforms3d + - matplotlib + - mediapy + - tyro + - msgpack + - websockets diff --git a/eval/simpler_env/policy/__init__.py b/eval/simpler_env/policy/__init__.py new file mode 100644 index 0000000..3dbcd28 --- /dev/null +++ b/eval/simpler_env/policy/__init__.py @@ -0,0 +1 @@ +"""Policies used by the SimplerEnv runners.""" diff --git a/eval/simpler_env/policy/model_server.py b/eval/simpler_env/policy/model_server.py new file mode 100644 index 0000000..b5fb743 --- /dev/null +++ b/eval/simpler_env/policy/model_server.py @@ -0,0 +1,234 @@ +"""StarVLA model-server adapter for the SimplerEnv WidowX benchmark.""" + +from __future__ import annotations + +import time +from collections import deque +from typing import Any + +import numpy as np + +from eval.libero.policy.model_server import ServerTiming +from robot_client.python.model_client import ModelClient, ModelResponse + + +DEFAULT_IMAGE_NAME = "image_0" +DEFAULT_IMAGE_SIZE = (224, 224) +DEFAULT_ACTION_ENSEMBLE_HORIZON = 7 +DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA = 0.1 + + +class AdaptiveEnsembler: + """StarVLA's cosine-similarity weighted temporal action ensemble.""" + + def __init__(self, horizon: int, alpha: float = DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA): + if horizon <= 0: + raise ValueError("action ensemble horizon must be positive") + self.horizon = int(horizon) + self.alpha = float(alpha) + self._history: deque[np.ndarray] = deque(maxlen=self.horizon) + + def reset(self) -> None: + self._history.clear() + + def ensemble_action(self, action_chunk: np.ndarray) -> np.ndarray: + chunk = np.asarray(action_chunk) + if not np.issubdtype(chunk.dtype, np.floating): + chunk = chunk.astype(np.float32) + if chunk.ndim not in (1, 2): + raise ValueError(f"expected a 1D action or 2D action chunk, got shape={chunk.shape}") + if chunk.ndim == 2 and chunk.shape[0] < min(len(self._history) + 1, self.horizon): + raise ValueError("action chunk is shorter than the active ensemble history") + + self._history.append(chunk) + count = len(self._history) + if chunk.ndim == 1: + current_predictions = np.stack(tuple(self._history)) + else: + current_predictions = np.stack( + [prediction[index] for index, prediction in zip(range(count - 1, -1, -1), self._history)] + ) + + reference = current_predictions[-1] + dot = np.sum(current_predictions * reference, axis=1) + norms = np.linalg.norm(current_predictions, axis=1) * np.linalg.norm(reference) + cosine = dot / (norms + 1e-7) + weights = np.exp(self.alpha * cosine) + weights /= weights.sum() + return np.sum(weights[:, None] * current_predictions, axis=0) + + +def resize_image_area(image: np.ndarray, image_size: tuple[int, int]) -> np.ndarray: + """Match the official StarVLA SimplerEnv client's OpenCV INTER_AREA resize.""" + + array = np.asarray(image) + if array.ndim != 3 or array.shape[2] != 3: + raise ValueError(f"expected an HWC RGB image, got shape={array.shape}") + if array.dtype != np.uint8: + raise ValueError(f"expected a uint8 RGB image, got dtype={array.dtype}") + width, height = (int(image_size[0]), int(image_size[1])) + if width <= 0 or height <= 0: + raise ValueError("image dimensions must be positive") + if array.shape[:2] == (height, width): + return np.ascontiguousarray(array) + try: + import cv2 + except ImportError as exc: + raise RuntimeError("opencv-python-headless is required to resize SimplerEnv observations") from exc + return cv2.resize(array, (width, height), interpolation=cv2.INTER_AREA) + + +def euler_xyz_to_axis_angle(rotation_delta: np.ndarray) -> np.ndarray: + """Use the same static-XYZ Euler convention as StarVLA's official adapter.""" + + roll, pitch, yaw = np.asarray(rotation_delta, dtype=np.float64).reshape(3) * 0.5 + sr, cr = np.sin(roll), np.cos(roll) + sp, cp = np.sin(pitch), np.cos(pitch) + sy, cy = np.sin(yaw), np.cos(yaw) + quaternion = np.asarray( + [ + cr * cp * cy + sr * sp * sy, + sr * cp * cy - cr * sp * sy, + cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, + ], + dtype=np.float64, + ) + quaternion /= np.linalg.norm(quaternion) + vector_norm = float(np.linalg.norm(quaternion[1:])) + if vector_norm <= 1e-12: + return np.zeros(3, dtype=np.float64) + angle = 2.0 * np.arccos(np.clip(quaternion[0], -1.0, 1.0)) + return quaternion[1:] * (angle / vector_norm) + + +class SimplerEnvModelServerPolicy: + """Closed-loop WidowX policy matching StarVLA's official SimplerEnv adapter.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 5555, + timeout: float | None = 120.0, + image_name: str = DEFAULT_IMAGE_NAME, + image_size: tuple[int, int] = DEFAULT_IMAGE_SIZE, + action_scale: float = 1.0, + action_ensemble: bool = True, + action_ensemble_horizon: int = DEFAULT_ACTION_ENSEMBLE_HORIZON, + adaptive_ensemble_alpha: float = DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA, + client: ModelClient | None = None, + ): + self.client = client or ModelClient(host=host, port=port, timeout=timeout) + self.image_name = str(image_name) + self.image_size = (int(image_size[0]), int(image_size[1])) + self.action_scale = float(action_scale) + self.action_ensembler = ( + AdaptiveEnsembler(action_ensemble_horizon, adaptive_ensemble_alpha) + if action_ensemble + else None + ) + self.task_description: str | None = None + self.predict_calls = 0 + self.timing_records: list[ServerTiming] = [] + self._action_shape: tuple[int, int] | None = None + + def health(self) -> str: + return self.client.health() + + def action_shape(self) -> tuple[int, int]: + if self._action_shape is None: + raise RuntimeError("model-server has not returned an action chunk") + return self._action_shape + + def _validate_response_actions(self, response: ModelResponse) -> np.ndarray: + shape = (int(response.chunk_size), int(response.action_dim)) + if shape[0] <= 0 or shape[1] <= 0: + raise RuntimeError(f"model-server returned an invalid action shape: {shape}") + # Protocol actions are FP32. Preserve that dtype through temporal + # ensembling to match StarVLA's official SimplerEnv client. + actions = np.asarray(response.actions, dtype=np.float32) + if actions.shape != shape: + raise RuntimeError( + "model-server returned an invalid action matrix: " + f"wire_shape={shape}, decoded_shape={actions.shape}" + ) + if not np.isfinite(actions).all(): + raise RuntimeError("model-server returned non-finite action values") + if shape[1] != 7: + raise RuntimeError(f"SimplerEnv WidowX requires 7D actions, got {shape}") + if self.action_ensembler is not None and shape[0] < self.action_ensembler.horizon: + raise RuntimeError( + f"action chunk {shape[0]} is shorter than ensemble horizon " + f"{self.action_ensembler.horizon}" + ) + if self._action_shape is not None and shape != self._action_shape: + raise RuntimeError( + f"model-server action shape changed from {self._action_shape} to {shape}" + ) + self._action_shape = shape + return actions + + def reset( + self, + task_description: str | None = None, + *, + reset_server: bool = True, + ) -> None: + self.task_description = task_description + if self.action_ensembler is not None: + self.action_ensembler.reset() + if reset_server: + self.client.reset() + + def build_observation(self, image: np.ndarray, task_description: str) -> dict[str, Any]: + resized = resize_image_area(image, self.image_size) + observation = { + "images": [{"name": self.image_name, "image": resized}], + "state": [], + "prompt": task_description, + } + return observation + + def predict_action_chunk(self, image: np.ndarray, task_description: str) -> ModelResponse: + request = self.build_observation(image, task_description) + started = time.perf_counter() + response = self.client.predict(request) + self._validate_response_actions(response) + self.timing_records.append( + ServerTiming( + roundtrip_ms=(time.perf_counter() - started) * 1000.0, + timings=response.timings, + ) + ) + self.predict_calls += 1 + return response + + def step( + self, image: np.ndarray, task_description: str | None = None + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + if task_description is not None and task_description != self.task_description: + self.reset(task_description, reset_server=False) + if self.task_description is None: + raise ValueError("task_description must be set before policy.step") + + response = self.predict_action_chunk(image, self.task_description) + actions = np.asarray(response.actions, dtype=np.float32) + selected = ( + self.action_ensembler.ensemble_action(actions) + if self.action_ensembler is not None + else actions[0] + ) + + raw_action = { + "world_vector": selected[:3].copy(), + "rotation_delta": selected[3:6].copy(), + "open_gripper": selected[6:7].copy(), + } + action = { + "world_vector": raw_action["world_vector"] * self.action_scale, + "rot_axangle": euler_xyz_to_axis_angle(raw_action["rotation_delta"]) * self.action_scale, + "gripper": 2.0 * (raw_action["open_gripper"] > 0.5).astype(np.float64) - 1.0, + "terminate_episode": np.asarray([0.0], dtype=np.float64), + } + return raw_action, action diff --git a/eval/simpler_env/runners/__init__.py b/eval/simpler_env/runners/__init__.py new file mode 100644 index 0000000..8145c73 --- /dev/null +++ b/eval/simpler_env/runners/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation runners.""" diff --git a/eval/simpler_env/runners/latency_starvla.py b/eval/simpler_env/runners/latency_starvla.py new file mode 100644 index 0000000..52972a7 --- /dev/null +++ b/eval/simpler_env/runners/latency_starvla.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import contextlib +import gc +import io +import json +import statistics +import subprocess +import sys +import tarfile +import tempfile +import time +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from eval.libero.utils.common import DEFAULT_RESULTS_DIR, timestamp, write_json + + +REPO_ROOT = Path(__file__).resolve().parents[3] +STARVLA_TOOLS = REPO_ROOT / "tools" / "hf2gguf" / "starvla" +if str(STARVLA_TOOLS) not in sys.path: + sys.path.insert(0, str(STARVLA_TOOLS)) + +from starvla_checkpoint import ( # noqa: E402 + DEFAULT_CATALOG, + get_qwen_asset, + get_variant, + load_catalog, + resolve_effective_config, +) + + +DEFAULT_PROMPT = "grab the block." +LEGACY_PI_REVISION = "e872a8579055f9332add8a2549b9fd5599e11510" +FULL_BF16_VARIANTS = {"oft", "qwen25_oft", "qwen25_pi"} +FRAMEWORK_CLASSES = { + "oft": ("starVLA.model.framework.VLM4A.QwenOFT", "Qwenvl_OFT"), + "groot": ("starVLA.model.framework.VLM4A.QwenGR00T", "Qwen_GR00T"), + "pi_v3": ("starVLA.model.framework.VLM4A.QwenPI_v3", "Qwen_PI_v3"), + "fast": ("starVLA.model.framework.VLM4A.QwenFast", "Qwenvl_Fast"), +} + + +def build_parser() -> argparse.ArgumentParser: + variants = tuple(load_catalog(DEFAULT_CATALOG)["variants"]) + parser = argparse.ArgumentParser(description="Benchmark an official StarVLA checkpoint with PyTorch.") + parser.add_argument("--variant", choices=variants, required=True) + parser.add_argument("--checkpoint-root", type=Path, default=REPO_ROOT / "ckpts" / "starvla") + parser.add_argument("--starvla-source", type=Path) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--compile-model", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--compile-mode", default="default") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--loops", type=int, default=20) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--image-height", type=int, default=224) + parser.add_argument("--image-width", type=int, default=224) + parser.add_argument("--output", type=Path) + return parser + + +def percentile(values: list[float], pct: float) -> float: + if len(values) == 1: + return values[0] + ordered = sorted(values) + pos = (len(ordered) - 1) * pct / 100.0 + lo = int(pos) + hi = min(lo + 1, len(ordered) - 1) + return ordered[lo] * (hi - pos) + ordered[hi] * (pos - lo) + + +def summarize(values: list[float]) -> dict[str, float | int]: + return { + "count": len(values), + "avg": statistics.fmean(values), + "min": min(values), + "p50": percentile(values, 50), + "p90": percentile(values, 90), + "p99": percentile(values, 99), + "max": max(values), + } + + +def sync(device: str) -> None: + if device.startswith("cuda"): + torch.cuda.synchronize(torch.device(device)) + + +def checkpoint_paths(checkpoint_root: Path, variant_name: str) -> dict[str, Any]: + catalog = load_catalog(DEFAULT_CATALOG) + variant = get_variant(catalog, variant_name) + _qwen_name, qwen = get_qwen_asset(catalog, variant) + policy_dir = checkpoint_root / "sources" / variant["directory"] / variant["revision"] + qwen_dir = checkpoint_root / "sources" / qwen["directory"] / qwen["revision"] + checkpoint = policy_dir / variant["checkpoint"]["path"] + for path in (policy_dir, qwen_dir, checkpoint): + if not path.exists(): + raise FileNotFoundError(path) + return { + "catalog": catalog, + "variant": variant, + "policy_dir": policy_dir.resolve(), + "qwen_dir": qwen_dir.resolve(), + "checkpoint": checkpoint.resolve(), + } + + +def verify_source(source: Path, catalog: Mapping[str, Any]) -> str: + revision = subprocess.run( + ["git", "-C", str(source), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + expected = catalog["source_revisions"]["starvla"] + if revision != expected: + raise RuntimeError(f"StarVLA source revision must be {expected}, got {revision}") + changes = subprocess.run( + ["git", "-C", str(source), "status", "--porcelain"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + if changes: + raise RuntimeError(f"StarVLA source checkout is not clean:\n{changes}") + return revision + + +@contextlib.contextmanager +def qwen_alias(qwen_dir: Path, backbone: str) -> Iterator[Path]: + name = "Qwen3-VL-4B-Instruct" if backbone == "qwen3_vl" else "Qwen2.5-VL-3B-Instruct" + with tempfile.TemporaryDirectory(prefix="starvla-latency-qwen-") as temporary: + alias = Path(temporary) / name + alias.symlink_to(qwen_dir, target_is_directory=True) + yield alias + + +@contextlib.contextmanager +def config_only_qwen(qwen_dir: Path, backbone: str) -> Iterator[None]: + import transformers + + model_class = ( + transformers.Qwen3VLForConditionalGeneration + if backbone == "qwen3_vl" + else transformers.Qwen2_5_VLForConditionalGeneration + ) + original = model_class.__dict__.get("from_pretrained") + + def from_config_only(_model_id: str, **_kwargs: Any) -> Any: + config = transformers.AutoConfig.from_pretrained(qwen_dir, local_files_only=True) + config._attn_implementation = "sdpa" + previous = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + with transformers.modeling_utils.no_init_weights(): + return model_class(config) + finally: + torch.set_default_dtype(previous) + + model_class.from_pretrained = staticmethod(from_config_only) + try: + yield + finally: + if original is None: + delattr(model_class, "from_pretrained") + else: + model_class.from_pretrained = original + + +def extract_legacy_source(source: Path) -> tempfile.TemporaryDirectory[str]: + archive = subprocess.run( + ["git", "-C", str(source), "archive", "--format=tar", LEGACY_PI_REVISION], + check=True, + stdout=subprocess.PIPE, + ).stdout + holder: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory(prefix="starvla-pi-") + with tarfile.open(fileobj=io.BytesIO(archive), mode="r:") as stream: + stream.extractall( + holder.name, + filter=lambda member, target: None + if member.issym() or member.islnk() + else tarfile.data_filter(member, target), + ) + return holder + + +def fast_config(policy_dir: Path) -> dict[str, Any]: + import yaml + + config = yaml.safe_load((policy_dir / "config.yaml").read_text(encoding="utf-8")) + config["framework"]["name"] = "QwenFast" + config["framework"]["action_model"]["action_model_type"] = "FAST" + config["framework"]["action_model"]["action_horizon"] = 16 + return config + + +def install_policy_dtype_bridge(framework: Any, framework_name: str) -> None: + if framework_name not in {"groot", "pi_v3"}: + return + action_model = framework.action_model + original_encoder = action_model.action_encoder.forward + original_dit = action_model.model.forward + + def encoder(actions: Any, timesteps: Any) -> Any: + return original_encoder(actions.float(), timesteps) + + def dit(*args: Any, **kwargs: Any) -> Any: + conditioning = kwargs.get("encoder_hidden_states", args[1] if len(args) > 1 else None) + if isinstance(conditioning, (list, tuple)): + conditioning = [value.float() for value in conditioning] + else: + conditioning = conditioning.float() + if "encoder_hidden_states" in kwargs: + kwargs["encoder_hidden_states"] = conditioning + else: + args = (args[0], conditioning, *args[2:]) + return original_dit(*args, **kwargs) + + action_model.action_encoder.forward = encoder + action_model.model.forward = dit + + +def load_framework(paths: Mapping[str, Any], source: Path, device: str) -> tuple[Any, Any]: + import importlib + + import yaml + + variant = paths["variant"] + variant_name = variant["_catalog_key"] + framework_name = variant["framework"] + runtime_source = source + holder = None + if variant_name == "qwen25_pi": + holder = extract_legacy_source(source) + runtime_source = Path(holder.name) + + sys.path.insert(0, str(runtime_source)) + try: + if variant_name == "qwen25_pi": + module_name, class_name = "starVLA.model.framework.QwenPI", "Qwen_PI" + config = yaml.safe_load((paths["policy_dir"] / "config.yaml").read_text(encoding="utf-8")) + else: + module_name, class_name = FRAMEWORK_CLASSES[framework_name] + config = ( + fast_config(paths["policy_dir"]) + if framework_name == "fast" + else resolve_effective_config(paths["policy_dir"], variant_name, variant) + ) + + from starVLA.model.framework import share_tools + + with qwen_alias(paths["qwen_dir"], variant["backbone"]) as alias: + config["framework"]["qwenvl"]["base_vlm"] = str(alias) + config["framework"]["qwenvl"]["attn_implementation"] = "sdpa" + cfg = share_tools.dict_to_namespace(config) + cfg.trainer.pretrained_checkpoint = None + module = importlib.import_module(module_name) + if framework_name == "fast": + from starVLA.model.modules.action_model.fast_ActionHeader import Fast_Action_Tokenizer + + codec = paths["catalog"]["shared_assets"]["fast_codec"] + codec_dir = ( + paths["policy_dir"].parents[1] / codec["directory"] / codec["revision"] + ) + module.get_action_model = lambda config=None: Fast_Action_Tokenizer(str(codec_dir)) + with config_only_qwen(paths["qwen_dir"], variant["backbone"]): + framework = getattr(module, class_name)(cfg) + + state = torch.load(paths["checkpoint"], map_location="cpu", mmap=True, weights_only=True) + framework.load_state_dict(state, strict=True) + del state + gc.collect() + framework.norm_stats = json.loads( + (paths["policy_dir"] / "dataset_statistics.json").read_text(encoding="utf-8") + ) + if variant_name in FULL_BF16_VARIANTS: + framework = framework.to(dtype=torch.bfloat16) + framework = framework.to(device).eval() + install_policy_dtype_bridge(framework, framework_name) + return framework, holder + except Exception: + if holder is not None: + holder.cleanup() + raise + finally: + if sys.path and sys.path[0] == str(runtime_source): + del sys.path[0] + + +def enable_compile(framework: Any, backbone: str, framework_name: str, mode: str) -> None: + qwen = framework.qwen_vl_interface + if framework_name == "fast": + qwen.model.forward = torch.compile(qwen.model.forward, mode=mode, fullgraph=False) + elif backbone == "qwen3_vl": + model = qwen.model.model + model.visual.forward = torch.compile(model.visual.forward, mode=mode, fullgraph=False) + for layer in model.language_model.layers: + layer.forward = torch.compile(layer.forward, mode=mode, fullgraph=False) + else: + qwen.forward = torch.compile(qwen.forward, mode=mode, fullgraph=False) + + if framework_name == "oft": + framework.action_model.predict_action = torch.compile( + framework.action_model.predict_action, mode=mode, fullgraph=False + ) + elif framework_name != "fast": + framework.action_model.model.forward = torch.compile( + framework.action_model.model.forward, mode=mode, fullgraph=False + ) + + +def unnormalize(normalized: Any, statistics: Mapping[str, Any]) -> np.ndarray: + profile_name = next(iter(statistics)) + stats = statistics[profile_name]["action"] + values = np.asarray(normalized, dtype=np.float32) + if values.shape != (1, 16, 7) or not np.isfinite(values).all(): + raise ValueError(f"StarVLA returned invalid normalized actions: {values.shape}") + q01 = np.asarray(stats["q01"], dtype=np.float32) + q99 = np.asarray(stats["q99"], dtype=np.float32) + mask = np.asarray(stats["mask"], dtype=np.bool_) + result = np.empty_like(values) + result[..., mask] = (values[..., mask] + 1.0) * 0.5 * (q99[mask] - q01[mask]) + q01[mask] + result[..., ~mask] = (values[..., ~mask] > 0.5).astype(np.float32) + return result + + +def predict(framework: Any, variant: str, image: Image.Image, prompt: str) -> Mapping[str, Any]: + if variant == "qwen25_pi": + return framework.predict_action(batch_images=[[image]], instructions=[prompt], state=None) + return framework.predict_action(examples=[{"image": [image], "lang": prompt}]) + + +def main() -> int: + args = build_parser().parse_args() + if args.warmup < 0 or args.loops <= 0: + raise ValueError("--warmup must be non-negative and --loops must be positive") + if not args.device.startswith("cuda") or not torch.cuda.is_available(): + raise RuntimeError("StarVLA latency currently requires CUDA") + + torch.manual_seed(args.seed) + np.random.seed(args.seed) + checkpoint_root = args.checkpoint_root.resolve() + source = (args.starvla_source or checkpoint_root / "source" / "starvla").resolve() + paths = checkpoint_paths(checkpoint_root, args.variant) + source_revision = verify_source(source, paths["catalog"]) + output = args.output or DEFAULT_RESULTS_DIR / f"starvla-policy-latency-{args.variant}-{timestamp()}.json" + + load_start = time.perf_counter() + framework, holder = load_framework(paths, source, args.device) + if args.compile_model: + enable_compile(framework, paths["variant"]["backbone"], paths["variant"]["framework"], args.compile_mode) + load_ms = (time.perf_counter() - load_start) * 1000.0 + + rng = np.random.default_rng(args.seed) + image = Image.fromarray( + rng.integers(0, 256, size=(args.image_height, args.image_width, 3), dtype=np.uint8), mode="RGB" + ) + rows: list[dict[str, float]] = [] + actions = None + print( + f"StarVLA latency: variant={args.variant} warmup={args.warmup} loops={args.loops} " + f"compile_model={args.compile_model} device={args.device}" + ) + for index in range(args.warmup + args.loops): + sync(args.device) + started = time.perf_counter() + output_value = predict(framework, args.variant, image, args.prompt) + sync(args.device) + policy_ms = (time.perf_counter() - started) * 1000.0 + + unnorm_started = time.perf_counter() + actions = unnormalize(output_value["normalized_actions"], framework.norm_stats) + unnormalize_ms = (time.perf_counter() - unnorm_started) * 1000.0 + total_ms = policy_ms + unnormalize_ms + if index >= args.warmup: + rows.append({"policy_ms": policy_ms, "unnormalize_ms": unnormalize_ms, "total_ms": total_ms}) + print( + f"iter={index} policy_ms={policy_ms:.3f} unnormalize_ms={unnormalize_ms:.3f} " + f"total_ms={total_ms:.3f}", + flush=True, + ) + + assert actions is not None + variant = paths["variant"] + payload = { + "runner": "starvla-policy-latency", + "variant": args.variant, + "framework": variant["framework"], + "backbone": variant["backbone"], + "checkpoint": {"repo_id": variant["repo_id"], "revision": variant["revision"]}, + "starvla_revision": source_revision, + "device": args.device, + "compile_model": args.compile_model, + "compile_mode": args.compile_mode if args.compile_model else None, + "warmup": args.warmup, + "loops": args.loops, + "load_ms": load_ms, + "action_shape": list(actions.shape), + "raw_input": {"image_shape_hwc": [args.image_height, args.image_width, 3], "prompt": args.prompt}, + "timing_ms": { + key: summarize([row[key] for row in rows]) + for key in ("policy_ms", "unnormalize_ms", "total_ms") + }, + "rows": rows, + } + write_json(output, payload) + print(f"wrote {output}") + print(json.dumps(payload["timing_ms"], indent=2)) + if holder is not None: + holder.cleanup() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/eval/simpler_env/runners/run_model_server.py b/eval/simpler_env/runners/run_model_server.py new file mode 100755 index 0000000..74dcd8c --- /dev/null +++ b/eval/simpler_env/runners/run_model_server.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import argparse +import time +from collections import defaultdict +from copy import copy +from pathlib import Path +from typing import Any + +import numpy as np + +from eval.libero.policy.model_server import ( + average_timing, + maybe_launch_server, + parse_server_env, + server_command, + stop_server, + timing_summary, +) +from eval.libero.utils.common import DEFAULT_RESULTS_DIR, aggregate_episodes, timestamp, write_json +from eval.simpler_env.policy.model_server import ( + DEFAULT_ACTION_ENSEMBLE_HORIZON, + DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA, + DEFAULT_IMAGE_NAME, + DEFAULT_IMAGE_SIZE, + SimplerEnvModelServerPolicy, +) +from eval.simpler_env.utils.environment import ( + BRIDGE_EPISODE_COUNT, + BRIDGE_SUITE, + BRIDGE_TASKS, + BridgeTask, + apply_runtime_env, + close_env, + language_instruction, + make_env, + observation_image, + parse_episode_ids, + parse_task_ids, + reset_env, + selected_tasks, + simpler_env_root, +) + + +def _first_bool(value: Any) -> bool: + array = np.asarray(value).reshape(-1) + if array.size != 1: + raise RuntimeError(f"expected one termination value, got shape={np.asarray(value).shape}") + return bool(array[0]) + + +def _first_float(value: Any) -> float: + array = np.asarray(value).reshape(-1) + if array.size != 1 or not np.isfinite(array[0]): + raise RuntimeError(f"expected one finite reward value, got {value!r}") + return float(array[0]) + + +def _write_video(path: Path, images: list[np.ndarray], fps: int) -> None: + try: + from simpler_env.utils.visualization import write_video + except ImportError as exc: + raise RuntimeError("failed to import SimplerEnv video writer") from exc + path.parent.mkdir(parents=True, exist_ok=True) + write_video(str(path), images, fps=fps) + + +def _command_with_noise_seed(command: list[str], seed: int) -> list[str]: + result = list(command) + for index, value in enumerate(result): + if value == "--noise-seed": + result[index + 1] = str(seed) + return result + if value.startswith("--noise-seed="): + result[index] = f"--noise-seed={seed}" + return result + return [*result, "--noise-seed", str(seed)] + + +def _launch_fresh_server(args: argparse.Namespace, policy: SimplerEnvModelServerPolicy): + try: + health = policy.health() + except OSError: + pass + else: + raise RuntimeError( + f"refusing to reuse model-server at {args.host}:{args.port}: {health}" + ) + process = maybe_launch_server(args, policy) + if process is None: + raise RuntimeError("model-server launch did not create a process") + return process + + +def model_record( + args: argparse.Namespace, policy: SimplerEnvModelServerPolicy +) -> dict[str, Any]: + chunk_size, action_dim = policy.action_shape() + return { + "model_type": args.expected_model_type, + "variant": args.variant, + "framework": args.expected_framework, + "checkpoint_revision": args.expected_checkpoint_revision, + "checkpoint_sha256": args.expected_checkpoint_sha256, + "qwen_revision": args.expected_qwen_revision, + "starvla_revision": args.expected_starvla_revision, + "chunk_size": chunk_size, + "action_dim": action_dim, + } + + +def aggregate_task_repeats(episodes: list[dict[str, Any]]) -> list[dict[str, Any]]: + groups: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) + for episode in episodes: + groups[(int(episode["task_id"]), int(episode["repeat"]))].append(episode) + return [ + { + "suite": BRIDGE_SUITE, + "task_id": task_id, + "repeat": repeat, + **aggregate_episodes(rows)["overall"], + } + for (task_id, repeat), rows in sorted(groups.items()) + ] + + +def run_episode( + env: Any, + policy: SimplerEnvModelServerPolicy, + task_spec: BridgeTask, + episode_id: int, + *, + repeat: int, + max_episode_steps: int, + camera_name: str | None, + video_path: Path | None, + video_fps: int, +) -> dict[str, Any]: + observation, _ = reset_env(env, task_spec, episode_id) + task = language_instruction(env) + if task != task_spec.instruction: + raise RuntimeError( + f"unexpected instruction for {task_spec.env_name}: {task!r}" + ) + policy.reset(task, reset_server=True) + image = observation_image(env, observation, camera_name) + frames = [image] if video_path is not None else [] + start_predict_calls = policy.predict_calls + start_timing_index = len(policy.timing_records) + started = time.perf_counter() + rewards: list[float] = [] + success = False + terminated = False + truncated = False + steps = 0 + + while steps < max_episode_steps and not truncated: + _, action = policy.step(image, task) + env_action = np.concatenate( + [action["world_vector"], action["rot_axangle"], action["gripper"]] + ) + observation, reward, terminated_value, truncated_value, _ = env.step(env_action) + terminated = _first_bool(terminated_value) + truncated = _first_bool(truncated_value) + success = terminated + rewards.append(_first_float(reward)) + steps += 1 + if terminated or truncated: + break + task = language_instruction(env) + image = observation_image(env, observation, camera_name) + if frames: + frames.append(image) + + if video_path is not None: + _write_video(video_path, frames, video_fps) + records = policy.timing_records[start_timing_index:] + return { + "episode": int(episode_id), + "repeat": int(repeat), + "task": task_spec.instruction, + "task_name": task_spec.name, + "env_name": task_spec.env_name, + "success": bool(success), + "terminated": bool(terminated), + "truncated": bool(truncated), + "sum_reward": float(sum(rewards)), + "max_reward": float(max(rewards) if rewards else 0.0), + "steps": steps, + "elapsed_s": time.perf_counter() - started, + "predict_calls": policy.predict_calls - start_predict_calls, + "server_timing_avg_ms": average_timing(records), + "video": str(video_path) if video_path is not None else None, + } + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Evaluate model-server on SimplerEnv Bridge") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=5555) + parser.add_argument("--launch-server", action="store_true") + parser.add_argument("--server-command", nargs=argparse.REMAINDER, help="must be last") + parser.add_argument("--server-env", action="append") + parser.add_argument("--server-wait-s", type=float, default=180.0) + parser.add_argument("--server-noise-seed-base", type=int, default=0) + parser.add_argument("--variant") + parser.add_argument("--expected-model-type") + parser.add_argument("--expected-checkpoint-revision") + parser.add_argument("--expected-checkpoint-sha256") + parser.add_argument("--expected-qwen-revision") + parser.add_argument("--expected-starvla-revision") + parser.add_argument("--expected-framework") + parser.add_argument("--task-ids", default="all") + parser.add_argument("--episode-ids", default="0:24") + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--max-episode-steps", type=int, default=120) + parser.add_argument("--control-freq", type=int, default=5) + parser.add_argument("--sim-freq", type=int, default=500) + parser.add_argument("--image-name", default=DEFAULT_IMAGE_NAME) + parser.add_argument("--image-size", type=int, nargs=2, default=list(DEFAULT_IMAGE_SIZE)) + parser.add_argument("--action-scale", type=float, default=1.0) + parser.add_argument("--no-action-ensemble", action="store_true") + parser.add_argument( + "--action-ensemble-horizon", type=int, default=DEFAULT_ACTION_ENSEMBLE_HORIZON + ) + parser.add_argument( + "--adaptive-ensemble-alpha", type=float, default=DEFAULT_ADAPTIVE_ENSEMBLE_ALPHA + ) + parser.add_argument("--camera-name") + parser.add_argument("--no-rgb-overlay", action="store_true") + parser.add_argument("--enable-raytracing", action="store_true") + parser.add_argument("--simpler-env-root", type=Path) + parser.add_argument("--record-video", action="store_true") + parser.add_argument("--video-dir", type=Path) + parser.add_argument("--video-fps", type=int, default=5) + parser.add_argument("--output", type=Path) + return parser.parse_args(argv) + + +def _validate_args(args: argparse.Namespace) -> tuple[list[int], list[int]]: + positive = ( + args.repeats, + args.max_episode_steps, + args.control_freq, + args.sim_freq, + args.video_fps, + args.action_ensemble_horizon, + ) + if any(value <= 0 for value in positive): + raise ValueError("repeat, episode, timing, and ensemble values must be positive") + if args.server_noise_seed_base < 0: + raise ValueError("--server-noise-seed-base must be non-negative") + return parse_task_ids(args.task_ids), parse_episode_ids(args.episode_ids) + + +def run(args: argparse.Namespace) -> dict[str, Any]: + task_ids, episode_ids = _validate_args(args) + output = args.output or DEFAULT_RESULTS_DIR / f"server-simpler-env-bridge-{timestamp()}.json" + video_dir = args.video_dir or output.with_suffix("").with_name(output.stem + "-videos") + apply_runtime_env() + root = simpler_env_root(args.simpler_env_root) + policy = SimplerEnvModelServerPolicy( + host=args.host, + port=args.port, + image_name=args.image_name, + image_size=tuple(args.image_size), + action_scale=args.action_scale, + action_ensemble=not args.no_action_ensemble, + action_ensemble_horizon=args.action_ensemble_horizon, + adaptive_ensemble_alpha=args.adaptive_ensemble_alpha, + ) + episodes: list[dict[str, Any]] = [] + launches: list[dict[str, Any]] = [] + recorded_model: dict[str, Any] | None = None + process = None + + try: + if not args.launch_server: + maybe_launch_server(args, policy) + base_command = server_command(args) if args.launch_server else [] + for task_spec in selected_tasks(task_ids): + for repeat in range(1, args.repeats + 1): + derived_seed = args.server_noise_seed_base + task_spec.task_id * args.repeats + repeat - 1 + noise_seed = None + if args.launch_server: + noise_seed = derived_seed + launch_args = copy(args) + launch_args.server_command = _command_with_noise_seed(base_command, noise_seed) + process = _launch_fresh_server(launch_args, policy) + launches.append( + { + "task_id": task_spec.task_id, + "repeat": repeat, + "noise_seed": noise_seed, + } + ) + for episode_id in episode_ids: + video_path = ( + video_dir + / f"task-{task_spec.task_id}-{task_spec.name}" + / f"repeat-{repeat:02d}-episode-{episode_id:02d}.mp4" + if args.record_video + else None + ) + env = make_env( + task_spec, + root=root, + control_freq=args.control_freq, + sim_freq=args.sim_freq, + max_episode_steps=args.max_episode_steps, + use_rgb_overlay=not args.no_rgb_overlay, + enable_raytracing=args.enable_raytracing, + ) + try: + result = run_episode( + env, + policy, + task_spec, + episode_id, + repeat=repeat, + max_episode_steps=args.max_episode_steps, + camera_name=args.camera_name, + video_path=video_path, + video_fps=args.video_fps, + ) + finally: + close_env(env) + result.update( + suite=BRIDGE_SUITE, + task_id=task_spec.task_id, + noise_seed=noise_seed, + ) + episodes.append(result) + print( + f"bridge[{task_spec.task_id}] repeat={repeat} episode={episode_id} " + f"success={result['success']} steps={result['steps']}" + ) + current_model = model_record(args, policy) + if recorded_model is None: + recorded_model = current_model + elif current_model != recorded_model: + raise RuntimeError("model action contract changed between repeats") + if process is not None: + stop_server(process, policy) + process = None + finally: + stop_server(process, policy) + + assert recorded_model is not None + full_coverage = ( + task_ids == [task.task_id for task in BRIDGE_TASKS] + and episode_ids == list(range(BRIDGE_EPISODE_COUNT)) + ) + payload = { + "runner": "model-server", + "benchmark": { + "name": "SimplerEnv WidowX Bridge", + "suite": BRIDGE_SUITE, + "coverage": "full" if full_coverage else "partial", + }, + "config": { + "task_ids": task_ids, + "episode_ids": episode_ids, + "repeats": args.repeats, + "max_episode_steps": args.max_episode_steps, + "control_freq": args.control_freq, + "sim_freq": args.sim_freq, + "host": args.host, + "port": args.port, + "server_command": base_command or None, + "server_env": parse_server_env(args.server_env), + "server_launches": launches, + "image_name": args.image_name, + "image_size": args.image_size, + "action_scale": args.action_scale, + "action_ensemble": not args.no_action_ensemble, + "action_ensemble_horizon": args.action_ensemble_horizon, + "adaptive_ensemble_alpha": args.adaptive_ensemble_alpha, + "rgb_overlay": not args.no_rgb_overlay, + "camera_name": args.camera_name, + "raytracing": args.enable_raytracing, + }, + "model": recorded_model, + "episodes": episodes, + "per_task_repeat": aggregate_task_repeats(episodes), + "timing_ms": timing_summary(policy.timing_records), + **aggregate_episodes(episodes), + } + write_json(output, payload) + print(f"wrote {output}") + print(f"overall: {payload['overall']}") + return payload + + +def main() -> int: + run(parse_args()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/eval/simpler_env/scripts/run_model_server.sh b/eval/simpler_env/scripts/run_model_server.sh new file mode 100755 index 0000000..1d3768a --- /dev/null +++ b/eval/simpler_env/scripts/run_model_server.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." >/dev/null 2>&1 && pwd)" +cd "${REPO_ROOT}" + +source "${REPO_ROOT}/tools/hf2gguf/starvla/starvla_variant_config.sh" +VARIANT="${VARIANT:-groot}" +load_starvla_variant "${VARIANT}" +GGUF_DIR="${GGUF_DIR:-ckpts/starvla/gguf/${VARIANT}}" +LLM_GGUF="${LLM_GGUF:-${GGUF_DIR}/qwen-${ARTIFACT_STEM}-bf16.gguf}" +MMPROJ_GGUF="${MMPROJ_GGUF:-${GGUF_DIR}/mmproj-${ARTIFACT_STEM}-bf16.gguf}" +if [[ "${VARIANT}" == "qwen25_fast" ]]; then + POLICY_GGUF="${POLICY_GGUF:-${GGUF_DIR}/policy-qwen25-fast.gguf}" +else + POLICY_GGUF="${POLICY_GGUF:-${GGUF_DIR}/starvla-${ARTIFACT_STEM}-policy-fp32.gguf}" +fi + +BACKEND="${BACKEND:-linux-cuda}" +case "${BACKEND}" in + linux-cuda) BUILD_DIR="${BUILD_DIR:-${REPO_ROOT}/build_cuda}" ;; + linux-cpu) BUILD_DIR="${BUILD_DIR:-${REPO_ROOT}/build}" ;; + *) echo "unsupported BACKEND=${BACKEND}; expected linux-cuda or linux-cpu" >&2; exit 2 ;; +esac +SERVER_BIN="${SERVER_BIN:-${BUILD_DIR}/bin/model-server}" +PYTHON_BIN="${PYTHON:-ckpts/simpler_env/.venv/bin/python}" +HOST="${HOST:-127.0.0.1}" +PORT="${PORT:-5555}" +NOISE_SEED_BASE="${NOISE_SEED_BASE:-1000}" + +for path in "${LLM_GGUF}" "${MMPROJ_GGUF}" "${POLICY_GGUF}"; do + if [[ ! -f "${path}" ]]; then + echo "missing GGUF: ${path}" >&2 + exit 2 + fi +done +if [[ ! -x "${SERVER_BIN}" ]]; then + echo "model-server was not found or is not executable: ${SERVER_BIN}" >&2 + exit 2 +fi +if [[ ! -x "${PYTHON_BIN}" ]]; then + echo "SimplerEnv Python was not found: ${PYTHON_BIN}" >&2 + echo "follow eval/simpler_env/README_ZH.md to create it" >&2 + exit 2 +fi + +eval_cmd=( + "${PYTHON_BIN}" -m eval.simpler_env.runners.run_model_server + --launch-server + --host "${HOST}" + --port "${PORT}" + --variant "${VARIANT}" + --expected-model-type "${MODEL_TYPE}" + --expected-checkpoint-revision "${CHECKPOINT_REVISION}" + --expected-checkpoint-sha256 "${CHECKPOINT_SHA256}" + --expected-qwen-revision "${QWEN_REVISION}" + --expected-starvla-revision "${STARVLA_REVISION}" + --expected-framework "${FRAMEWORK}" + --server-noise-seed-base "${NOISE_SEED_BASE}" +) +[[ -n "${TASK_IDS:-}" ]] && eval_cmd+=(--task-ids "${TASK_IDS}") +[[ -n "${EPISODE_IDS:-}" ]] && eval_cmd+=(--episode-ids "${EPISODE_IDS}") +[[ -n "${REPEATS:-}" ]] && eval_cmd+=(--repeats "${REPEATS}") +[[ -n "${OUTPUT:-}" ]] && eval_cmd+=(--output "${OUTPUT}") +[[ -n "${SIMPLER_ENV_ROOT:-}" ]] && eval_cmd+=(--simpler-env-root "${SIMPLER_ENV_ROOT}") +[[ "${RECORD_VIDEO:-0}" == "1" ]] && eval_cmd+=(--record-video) +eval_cmd+=("$@") +eval_cmd+=( + --server-command + "${SERVER_BIN}" + --model-type "${MODEL_TYPE}" + --policy "${POLICY_GGUF}" + --llm "${LLM_GGUF}" + --mmproj "${MMPROJ_GGUF}" + --host "${HOST}" + --port "${PORT}" + --noise-seed "${NOISE_SEED_BASE}" +) + +exec "${eval_cmd[@]}" diff --git a/eval/simpler_env/utils/__init__.py b/eval/simpler_env/utils/__init__.py new file mode 100644 index 0000000..ea81195 --- /dev/null +++ b/eval/simpler_env/utils/__init__.py @@ -0,0 +1 @@ +"""SimplerEnv evaluation helpers.""" diff --git a/eval/simpler_env/utils/environment.py b/eval/simpler_env/utils/environment.py new file mode 100644 index 0000000..4978636 --- /dev/null +++ b/eval/simpler_env/utils/environment.py @@ -0,0 +1,230 @@ +"""Official StarVLA WidowX Bridge task and environment configuration.""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + + +BRIDGE_SUITE = "simpler_env_widowx_bridge" +BRIDGE_EPISODE_COUNT = 24 +BRIDGE_OFFICIAL_REPEATS = 4 +BRIDGE_CONTROL_MODE = "arm_pd_ee_target_delta_pose_align2_gripper_pd_joint_pos" + + +@dataclass(frozen=True) +class BridgeTask: + task_id: int + name: str + env_name: str + instruction: str + scene_name: str + robot: str + overlay_filename: str + robot_init_x: float + robot_init_y: float + + +BRIDGE_TASKS = ( + BridgeTask( + 0, + "stack_green_cube_on_yellow_cube", + "StackGreenCubeOnYellowCubeBakedTexInScene-v0", + "stack the green block on the yellow block", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 1, + "put_carrot_on_plate", + "PutCarrotOnPlateInScene-v0", + "put carrot on plate", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 2, + "put_spoon_on_table_cloth", + "PutSpoonOnTableClothInScene-v0", + "put the spoon on the towel", + "bridge_table_1_v1", + "widowx", + "bridge_real_eval_1.png", + 0.147, + 0.028, + ), + BridgeTask( + 3, + "put_eggplant_in_basket", + "PutEggplantInBasketScene-v0", + "put eggplant into yellow basket", + "bridge_table_1_v2", + "widowx_sink_camera_setup", + "bridge_sink.png", + 0.127, + 0.060, + ), +) + + +def parse_task_ids(value: str | None) -> list[int]: + if value is None or value.strip().lower() in {"", "all"}: + return [task.task_id for task in BRIDGE_TASKS] + text = value.strip() + decoded = json.loads(text) if text.startswith("[") else text.split(",") + if not isinstance(decoded, list): + raise ValueError("--task-ids must be 'all', a comma list, or a JSON list") + task_ids = [int(item) for item in decoded] + known = {task.task_id for task in BRIDGE_TASKS} + if len(set(task_ids)) != len(task_ids) or any(task_id not in known for task_id in task_ids): + raise ValueError(f"--task-ids must contain unique values from {sorted(known)}") + return task_ids + + +def selected_tasks(task_ids: list[int]) -> list[BridgeTask]: + by_id = {task.task_id: task for task in BRIDGE_TASKS} + return [by_id[task_id] for task_id in task_ids] + + +def parse_episode_ids(value: str | None) -> list[int]: + text = (value or "0:24").strip() + if ":" in text and not text.startswith("["): + fields = text.split(":") + if len(fields) not in (2, 3): + raise ValueError("--episode-ids range must be START:STOP or START:STOP:STEP") + start, stop = int(fields[0]), int(fields[1]) + step = int(fields[2]) if len(fields) == 3 else 1 + if step <= 0: + raise ValueError("--episode-ids range step must be positive") + episode_ids = list(range(start, stop, step)) + else: + decoded = json.loads(text) if text.startswith("[") else text.split(",") + if not isinstance(decoded, list): + raise ValueError("--episode-ids must be a comma list, JSON list, or range") + episode_ids = [int(item) for item in decoded if str(item).strip()] + if not episode_ids: + raise ValueError("--episode-ids must select at least one episode") + if len(set(episode_ids)) != len(episode_ids): + raise ValueError("--episode-ids must not contain duplicates") + if any(episode_id < 0 or episode_id >= BRIDGE_EPISODE_COUNT for episode_id in episode_ids): + raise ValueError(f"Bridge episode ids must be in [0, {BRIDGE_EPISODE_COUNT})") + return episode_ids + + +def apply_runtime_env() -> None: + os.environ["DISPLAY"] = "" + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + + +def simpler_env_root(explicit_root: Path | None = None) -> Path: + try: + import mani_skill2_real2sim + import simpler_env + except ImportError as exc: + raise RuntimeError( + "simpler_env is not installed; follow eval/simpler_env/README_ZH.md" + ) from exc + + if explicit_root is not None: + root = explicit_root.expanduser().resolve() + else: + root = Path(simpler_env.__file__).resolve().parent.parent + if not (root / "ManiSkill2_real2sim").exists(): + raise RuntimeError(f"invalid SimplerEnv root (ManiSkill2_real2sim missing): {root}") + installed_simpler_root = Path(simpler_env.__file__).resolve().parent.parent + installed_maniskill_root = Path(mani_skill2_real2sim.__file__).resolve().parent.parent + expected_maniskill_root = (root / "ManiSkill2_real2sim").resolve() + if installed_simpler_root != root: + raise RuntimeError( + f"installed simpler_env comes from {installed_simpler_root}, expected {root}" + ) + if installed_maniskill_root != expected_maniskill_root: + raise RuntimeError( + "installed mani_skill2_real2sim comes from " + f"{installed_maniskill_root}, expected {expected_maniskill_root}" + ) + return root + + +def overlay_path(root: Path, task: BridgeTask) -> Path: + path = root / "ManiSkill2_real2sim" / "data" / "real_inpainting" / task.overlay_filename + if not path.is_file(): + raise RuntimeError(f"official Bridge RGB overlay is missing: {path}") + return path + + +def make_env( + task: BridgeTask, + *, + root: Path, + control_freq: int = 5, + sim_freq: int = 500, + max_episode_steps: int = 120, + use_rgb_overlay: bool = True, + enable_raytracing: bool = False, +) -> Any: + try: + from simpler_env.utils.env.env_builder import build_maniskill2_env + except ImportError as exc: + raise RuntimeError("failed to import the installed SimplerEnv environment builder") from exc + + additional: dict[str, Any] = {"shader_dir": "rt"} if enable_raytracing else {} + return build_maniskill2_env( + task.env_name, + **additional, + obs_mode="rgbd", + robot=task.robot, + sim_freq=int(sim_freq), + control_mode=BRIDGE_CONTROL_MODE, + control_freq=int(control_freq), + max_episode_steps=int(max_episode_steps), + scene_name=task.scene_name, + camera_cfgs={"add_segmentation": True}, + rgb_overlay_path=str(overlay_path(root, task)) if use_rgb_overlay else None, + ) + + +def reset_env(env: Any, task: BridgeTask, episode_id: int) -> tuple[Any, Any]: + options = { + "robot_init_options": { + "init_xy": np.asarray([task.robot_init_x, task.robot_init_y], dtype=np.float64), + "init_rot_quat": np.asarray([0.0, 0.0, 0.0, 1.0], dtype=np.float64), + }, + "obj_init_options": {"episode_id": int(episode_id)}, + } + return env.reset(options=options) + + +def observation_image(env: Any, observation: dict[str, Any], camera_name: str | None = None) -> np.ndarray: + try: + from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict + except ImportError as exc: + raise RuntimeError("failed to import SimplerEnv observation helpers") from exc + image = np.asarray(get_image_from_maniskill2_obs_dict(env, observation, camera_name=camera_name)) + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[2] != 3: + raise RuntimeError(f"SimplerEnv returned an invalid RGB observation: shape={image.shape}, dtype={image.dtype}") + return image + + +def language_instruction(env: Any) -> str: + instruction = str(env.get_language_instruction()) + if not instruction: + raise RuntimeError("SimplerEnv returned an empty language instruction") + return instruction + + +def close_env(env: Any) -> None: + close = getattr(env, "close", None) + if callable(close): + close() diff --git a/patches/llama.cpp/0001-qwen3vl-vision-parity.patch b/patches/llama.cpp/0001-qwen3vl-vision-parity.patch new file mode 100644 index 0000000..a868059 --- /dev/null +++ b/patches/llama.cpp/0001-qwen3vl-vision-parity.patch @@ -0,0 +1,36 @@ +diff --git a/tools/mtmd/models/qwen3vl.cpp b/tools/mtmd/models/qwen3vl.cpp +index fa1100d..5119df6 100644 +--- a/tools/mtmd/models/qwen3vl.cpp ++++ b/tools/mtmd/models/qwen3vl.cpp +@@ -43,8 +43,11 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + cb(inp, "patch_bias", -1); + } + +- // calculate absolute position embedding and apply +- ggml_tensor * learned_pos_embd = resize_position_embeddings(); ++ // Qwen3-VL constructs interpolation coordinates with torch.linspace(0, ++ // num_grid_per_side - 1, size), which is bilinear align_corners=True ++ // without antialiasing. ++ ggml_tensor * learned_pos_embd = resize_position_embeddings( ++ GGML_SCALE_MODE_BILINEAR | GGML_SCALE_FLAG_ALIGN_CORNERS); + learned_pos_embd = ggml_cont_4d( + ctx0, learned_pos_embd, + n_embd * 2, n_patches_x / 2, n_patches_y, batch_size); +@@ -154,7 +157,7 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + layer.deepstack_fc1_w, layer.deepstack_fc1_b, + nullptr, nullptr, + layer.deepstack_fc2_w, layer.deepstack_fc2_b, +- ffn_op_type::FFN_GELU, il); ++ ffn_op_type::FFN_GELU_ERF, il); + + if(!deepstack_features) { + deepstack_features = feat; +@@ -180,7 +183,7 @@ ggml_cgraph * clip_graph_qwen3vl::build() { + model.mm_0_w, model.mm_0_b, + nullptr, nullptr, + model.mm_1_w, model.mm_1_b, +- ffn_op_type::FFN_GELU, -1); ++ ffn_op_type::FFN_GELU_ERF, -1); + + if (deepstack_features) { + embeddings = ggml_concat(ctx0, embeddings, deepstack_features, 0); diff --git a/patches/llama.cpp/0002-per-context-native-graph-control.patch b/patches/llama.cpp/0002-per-context-native-graph-control.patch new file mode 100644 index 0000000..3b3a992 --- /dev/null +++ b/patches/llama.cpp/0002-per-context-native-graph-control.patch @@ -0,0 +1,267 @@ +diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h +index d0c7e5a..d86d0b3 100644 +--- a/ggml/include/ggml-backend.h ++++ b/ggml/include/ggml-backend.h +@@ -215,6 +215,8 @@ extern "C" { + typedef ggml_backend_buffer_type_t * (*ggml_backend_dev_get_extra_bufts_t)(ggml_backend_dev_t device); + // Set the abort callback for the backend + typedef void (*ggml_backend_set_abort_callback_t)(ggml_backend_t backend, ggml_abort_callback abort_callback, void * abort_callback_data); ++ // Enable or disable native graph capture/cache for one backend instance. ++ typedef void (*ggml_backend_set_native_graphs_enabled_t)(ggml_backend_t backend, bool enabled); + // Get a list of feature flags supported by the backend (returns a NULL-terminated array) + struct ggml_backend_feature { + const char * name; +diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh +index 1081750..fb7c279 100644 +--- a/ggml/src/ggml-cuda/common.cuh ++++ b/ggml/src/ggml-cuda/common.cuh +@@ -1373,6 +1373,8 @@ struct ggml_backend_cuda_context { + int curr_stream_no = 0; + + #ifdef USE_CUDA_GRAPH ++ bool cuda_graphs_enabled = true; ++ + // Map from first_node_ptr to cuda_graph - allows multiple graphs per context + // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) + std::unordered_map> cuda_graphs; +@@ -1405,6 +1407,9 @@ struct ggml_backend_cuda_context { + // Check if any CUDA graph is enabled for this context (used by kernels that need to know + // if graphs are in use without having access to the specific graph key) + bool any_cuda_graph_enabled() const { ++ if (!cuda_graphs_enabled) { ++ return false; ++ } + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->is_enabled()) { + return true; +@@ -1415,6 +1420,9 @@ struct ggml_backend_cuda_context { + + // Check if any CUDA graph has an instance for this context + bool any_cuda_graph_has_instance() const { ++ if (!cuda_graphs_enabled) { ++ return false; ++ } + for (const auto & [key, graph] : cuda_graphs) { + if (graph && graph->instance != nullptr) { + return true; +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index 8d21b22..42f7838 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -3085,6 +3085,23 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { + GGML_UNUSED(backend); + } + ++static void ggml_backend_cuda_set_native_graphs_enabled(ggml_backend_t backend, bool enabled) { ++ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ++ ++#ifdef USE_CUDA_GRAPH ++ if (cuda_ctx->cuda_graphs_enabled == enabled) { ++ return; ++ } ++ ++ ggml_backend_cuda_synchronize(backend); ++ cuda_ctx->cuda_graphs.clear(); ++ cuda_ctx->cuda_graphs_enabled = enabled; ++#else ++ GGML_UNUSED(cuda_ctx); ++ GGML_UNUSED(enabled); ++#endif ++} ++ + #ifdef USE_CUDA_GRAPH + static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { + +@@ -4202,8 +4219,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud + } + + #ifdef USE_CUDA_GRAPH +- ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture ++ ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->graph != nullptr) { + CUDA_CHECK(cudaGraphDestroy(graph->graph)); + graph->graph = nullptr; +@@ -4240,6 +4257,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud + + #ifdef USE_CUDA_GRAPH + static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { ++ if (!cuda_ctx->cuda_graphs_enabled) { ++ return false; ++ } ++ + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + if (graph->graph == nullptr) { +@@ -4267,10 +4288,8 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, + #ifdef USE_CUDA_GRAPH + graph_key = ggml_cuda_graph_get_key(cgraph); + +- ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); +- +- ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); +- if (graph->is_enabled()) { ++ if (ggml_cuda_graph_set_enabled(cuda_ctx, graph_key)) { ++ ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); + if (graph_compatible) { + const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); +@@ -5400,6 +5419,9 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cuda_get_features; + } ++ if (strcmp(name, "ggml_backend_set_native_graphs_enabled") == 0) { ++ return (void *)ggml_backend_cuda_set_native_graphs_enabled; ++ } + return nullptr; + } + +diff --git a/include/llama.h b/include/llama.h +index 2ea2267..aa6c656 100644 +--- a/include/llama.h ++++ b/include/llama.h +@@ -969,6 +969,10 @@ extern "C" { + // Set abort callback + LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); + ++ // Enable or disable native graph capture/cache for each context backend ++ // that exposes this optional capability. Direct graph computation remains enabled. ++ LLAMA_API void llama_set_backend_native_graphs_enabled(struct llama_context * ctx, bool enabled); ++ + // Wait until all computations are finished + // This is automatically done when using one of the functions below to obtain the computation results + // and is not necessary to call it explicitly in most cases +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index 71a5939..a705255 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -1031,6 +1031,21 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void + } + } + ++void llama_context::set_backend_native_graphs_enabled(bool enabled) { ++ for (auto & backend : backends) { ++ auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get())); ++ if (reg == nullptr) { ++ continue; ++ } ++ auto * set_enabled = reinterpret_cast( ++ ggml_backend_reg_get_proc_address( ++ reg, "ggml_backend_set_native_graphs_enabled")); ++ if (set_enabled != nullptr) { ++ set_enabled(backend.get(), enabled); ++ } ++ } ++} ++ + void llama_context::set_embeddings(bool value) { + LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); + +@@ -3365,6 +3380,10 @@ void llama_set_abort_callback(llama_context * ctx, bool (*abort_callback)(void * + ctx->set_abort_callback(abort_callback, abort_callback_data); + } + ++void llama_set_backend_native_graphs_enabled(llama_context * ctx, bool enabled) { ++ ctx->set_backend_native_graphs_enabled(enabled); ++} ++ + void llama_set_embeddings(llama_context * ctx, bool embeddings) { + ctx->set_embeddings(embeddings); + } +diff --git a/src/llama-context.h b/src/llama-context.h +index 92d1b0c..d1354a4 100644 +--- a/src/llama-context.h ++++ b/src/llama-context.h +@@ -105,6 +105,7 @@ struct llama_context { + void set_n_threads(int32_t n_threads, int32_t n_threads_batch); + + void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data); ++ void set_backend_native_graphs_enabled(bool enabled); + + void set_embeddings (bool value); + void set_causal_attn(bool value); +diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp +index 513b94f..4e63c93 100644 +--- a/tools/mtmd/clip.cpp ++++ b/tools/mtmd/clip.cpp +@@ -2973,6 +2973,26 @@ void clip_free(clip_ctx * ctx) { + delete ctx; + } + ++void clip_set_backend_native_graphs_enabled(clip_ctx * ctx, bool enabled) { ++ if (ctx == nullptr) { ++ return; ++ } ++ ++ for (ggml_backend_t backend : ctx->backend_ptrs) { ++ ggml_backend_dev_t dev = ggml_backend_get_device(backend); ++ ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr; ++ if (reg == nullptr) { ++ continue; ++ } ++ auto * set_enabled = reinterpret_cast( ++ ggml_backend_reg_get_proc_address( ++ reg, "ggml_backend_set_native_graphs_enabled")); ++ if (set_enabled != nullptr) { ++ set_enabled(backend, enabled); ++ } ++ } ++} ++ + // deprecated + size_t clip_embd_nbytes(const struct clip_ctx * ctx) { + const int32_t nx = ctx->model.hparams.image_size; +diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h +index a859b38..d63d43d 100644 +--- a/tools/mtmd/clip.h ++++ b/tools/mtmd/clip.h +@@ -51,6 +51,10 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params + + void clip_free(struct clip_ctx * ctx); + ++// Enable or disable native graph capture/cache for each CLIP backend that ++// exposes this optional capability. Direct graph computation remains enabled. ++void clip_set_backend_native_graphs_enabled(struct clip_ctx * ctx, bool enabled); ++ + size_t clip_embd_nbytes(const struct clip_ctx * ctx); + size_t clip_embd_nbytes_by_img(const struct clip_ctx * ctx, int img_w, int img_h); + +diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp +index 87da687..5fa1bc3 100644 +--- a/tools/mtmd/mtmd.cpp ++++ b/tools/mtmd/mtmd.cpp +@@ -628,6 +628,19 @@ void mtmd_free(mtmd_context * ctx) { + delete ctx; + } + ++void mtmd_set_backend_native_graphs_enabled(mtmd_context * ctx, bool enabled) { ++ if (ctx == nullptr) { ++ return; ++ } ++ ++ if (ctx->ctx_v != nullptr) { ++ clip_set_backend_native_graphs_enabled(ctx->ctx_v, enabled); ++ } ++ if (ctx->ctx_a != nullptr) { ++ clip_set_backend_native_graphs_enabled(ctx->ctx_a, enabled); ++ } ++} ++ + struct mtmd_tokenizer { + mtmd_context * ctx; + std::vector bitmaps; +diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h +index e364174..7cf3af7 100644 +--- a/tools/mtmd/mtmd.h ++++ b/tools/mtmd/mtmd.h +@@ -110,6 +110,10 @@ MTMD_API mtmd_context * mtmd_init_from_file(const char * mmproj_fname, + + MTMD_API void mtmd_free(mtmd_context * ctx); + ++// Enable or disable native graph capture/cache for each media backend that ++// exposes this optional capability. Direct graph computation remains enabled. ++MTMD_API void mtmd_set_backend_native_graphs_enabled(mtmd_context * ctx, bool enabled); ++ + // whether we need to set non-causal mask before llama_decode + // if chunk is nullptr, we assume the default case where chunk is an image chunk + MTMD_API bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk * chunk); diff --git a/patches/llama.cpp/README.md b/patches/llama.cpp/README.md new file mode 100644 index 0000000..1a17011 --- /dev/null +++ b/patches/llama.cpp/README.md @@ -0,0 +1,31 @@ +# llama.cpp patches + +The project pins `third_party/llama.cpp` at commit +`3e941b813b1acbbf06c2203a94ceb33d84748c1e`. The repository applies two +changes that are not available through that revision's public APIs: + +1. `0001-qwen3vl-vision-parity.patch` uses the position interpolation and exact + GELU operations from the Qwen3-VL implementation used by StarVLA. +2. `0002-per-context-native-graph-control.patch` adds an optional backend API to + disable CUDA graph capture for the text and vision contexts owned by one + StarVLA instance. This avoids retained CUDA graphs growing memory use during + long runs without changing the setting for other llama.cpp users. + +Apply the repository patch set after initializing submodules and before building: + +```bash +./tools/apply_patches.sh +``` + +The command verifies the exact llama.cpp revision and refuses a dirty or +partially patched checkout. It is safe to run again after a complete apply. + +Inspect or remove the overlay with: + +```bash +./tools/apply_patches.sh --check +./tools/apply_patches.sh --revert +``` + +The parent repository commits only these patch assets. It does not advance or +commit a forked llama.cpp gitlink. diff --git a/robot_client/cpp/model_client.cpp b/robot_client/cpp/model_client.cpp index 5496144..3ad82d8 100644 --- a/robot_client/cpp/model_client.cpp +++ b/robot_client/cpp/model_client.cpp @@ -55,8 +55,9 @@ bool make_predict_request(const ModelObservation & obs, proto::predict_request & return false; } - req.task = obs.prompt; - req.state = obs.state; + req.task = obs.prompt; + req.state = obs.state; + req.initial_noise = obs.initial_noise; req.images.clear(); req.images.reserve(obs.images.size()); diff --git a/robot_client/cpp/model_client.h b/robot_client/cpp/model_client.h index a6b7730..ea84f47 100644 --- a/robot_client/cpp/model_client.h +++ b/robot_client/cpp/model_client.h @@ -21,6 +21,7 @@ struct ModelImage { struct ModelObservation { std::vector images; std::vector state; + std::vector initial_noise; std::string prompt = "grab the block."; }; diff --git a/robot_client/python/model_client.py b/robot_client/python/model_client.py index 06e3791..5c79423 100644 --- a/robot_client/python/model_client.py +++ b/robot_client/python/model_client.py @@ -9,7 +9,7 @@ MAGIC = 0x414C5653 -VERSION = 3 +VERSION = 4 HEADER_SIZE = 32 OP_HEALTH = 1 @@ -21,8 +21,8 @@ IMAGE_RAW_RGB_U8 = 1 HEADER = struct.Struct(" tuple[int, int, int, bytes]: def encode_predict_observation(observation: dict[str, Any]) -> bytes: images = observation["images"] state = state_to_list(observation["state"]) + initial_noise = state_to_list(observation.get("initial_noise")) prompt = str(observation["prompt"]) if not images: raise ValueError("observation.images must contain at least one image") @@ -88,13 +89,14 @@ def encode_predict_observation(observation: dict[str, Any]) -> bytes: encoded_images.append((name, rgb, width, height, stride)) payload = bytearray() - payload += PREDICT_REQ_V2_FIXED.pack( + payload += PREDICT_REQ_FIXED.pack( len(encoded_images), len(state), + len(initial_noise), len(prompt_bytes), ) for name, rgb, width, height, stride in encoded_images: - payload += PREDICT_REQ_V2_IMAGE.pack( + payload += PREDICT_REQ_IMAGE.pack( IMAGE_RAW_RGB_U8, len(name), width, @@ -105,6 +107,8 @@ def encode_predict_observation(observation: dict[str, Any]) -> bytes: ) for value in state: payload += struct.pack(" @@ -26,6 +27,7 @@ struct server_args { std::string action_decoder_path; std::string state_proj_path; std::string action_expert_path; + std::string policy_path; std::string task = "grab the block."; std::string host = "127.0.0.1"; int port = 5555; @@ -37,18 +39,6 @@ struct server_args { int verbosity = 0; }; -static bool parse_model_type(const std::string & value, robotcpp::model_type & out) { - if (value == "smolvla") { - out = robotcpp::model_type::smolvla; - return true; - } - if (value == "pi0") { - out = robotcpp::model_type::pi0; - return true; - } - return false; -} - static bool parse_noise_mode(const std::string & value, int & out_mode) { if (value == "gaussian") { out_mode = SMOLVLA_NOISE_MODE_GAUSSIAN; @@ -74,9 +64,11 @@ static void print_usage(const char * prog) { " [options]\n" " %s --model-type pi0 --vit --mmproj --llm --tokenizer --state-gguf " " --action-decoder [options]\n" + " %s --model-type starvla --llm --mmproj --policy [options]\n" "\n" "Common options:\n" - " --model-type Model type (default: smolvla)\n" + " --model-type smolvla|pi0|starvla\n" + " (default: smolvla)\n" "\n" "SmolVLA options:\n" " --llm LLM GGUF path\n" @@ -93,6 +85,11 @@ static void print_usage(const char * prog) { " --state-gguf State projector GGUF path\n" " --action-decoder Action decoder GGUF path\n" "\n" + "StarVLA options:\n" + " --policy StarVLA policy GGUF path (required)\n" + " --llm Qwen text GGUF path (required)\n" + " --mmproj Qwen vision GGUF path (required)\n" + "\n" "Runtime options:\n" " --host Listen host (default: 127.0.0.1)\n" " --port Listen port (default: 5555)\n" @@ -103,7 +100,7 @@ static void print_usage(const char * prog) { " --noise-seed RNG seed, <0 means auto (default: -1)\n" " --verbosity Log verbosity (default: 0)\n" " -h, --help Show this help\n", - prog, prog); + prog, prog, prog); } // TODO: may need to be cleaned up and optimized @@ -116,10 +113,12 @@ static bool parse_args(int argc, char ** argv, server_args & args) { } else if (arg == "--llm" && i + 1 < argc) { args.llm_path = argv[++i]; } else if (arg == "--model-type" && i + 1 < argc) { - if (!parse_model_type(argv[++i], args.model_type)) { + if (!robotcpp::parse_model_type(argv[++i], args.model_type)) { std::fprintf(stderr, "Error: unsupported model type '%s'\n", argv[i]); return false; } + } else if (arg == "--policy" && i + 1 < argc) { + args.policy_path = argv[++i]; } else if (arg == "--mmproj" && i + 1 < argc) { args.mmproj_path = argv[++i]; } else if (arg == "--vit" && i + 1 < argc) { @@ -139,22 +138,46 @@ static bool parse_args(int argc, char ** argv, server_args & args) { } else if (arg == "--host" && i + 1 < argc) { args.host = argv[++i]; } else if (arg == "--port" && i + 1 < argc) { - args.port = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.port)) { + std::fprintf(stderr, "Error: invalid --port value '%s'\n", value); + return false; + } } else if (arg == "--threads" && i + 1 < argc) { - args.threads = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.threads)) { + std::fprintf(stderr, "Error: invalid --threads value '%s'\n", value); + return false; + } } else if (arg == "--n-batch" && i + 1 < argc) { - args.n_batch = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_batch)) { + std::fprintf(stderr, "Error: invalid --n-batch value '%s'\n", value); + return false; + } } else if (arg == "--n-ctx" && i + 1 < argc) { - args.n_ctx = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_ctx)) { + std::fprintf(stderr, "Error: invalid --n-ctx value '%s'\n", value); + return false; + } } else if (arg == "--noise-mode" && i + 1 < argc) { if (!parse_noise_mode(argv[++i], args.noise_mode)) { std::fprintf(stderr, "Error: invalid noise mode '%s'\n", argv[i]); return false; } } else if (arg == "--noise-seed" && i + 1 < argc) { - args.noise_seed = (int64_t)std::atoll(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.noise_seed)) { + std::fprintf(stderr, "Error: invalid --noise-seed value '%s'\n", value); + return false; + } } else if (arg == "--verbosity" && i + 1 < argc) { - args.verbosity = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.verbosity)) { + std::fprintf(stderr, "Error: invalid --verbosity value '%s'\n", value); + return false; + } } else { std::fprintf(stderr, "Error: unknown argument '%s'\n", arg.c_str()); return false; @@ -168,15 +191,36 @@ static bool parse_args(int argc, char ** argv, server_args & args) { std::fprintf(stderr, "Error: model-server only listens on 127.0.0.1 in this phase\n"); return false; } + if (args.threads < 0 || args.n_batch <= 0 || args.n_ctx <= 0 || args.verbosity < 0) { + std::fprintf(stderr, + "Error: --threads/--verbosity must be non-negative and --n-batch/--n-ctx must be positive\n"); + return false; + } + if (robotcpp::is_starvla_model_type(args.model_type) && args.noise_mode != SMOLVLA_NOISE_MODE_GAUSSIAN) { + std::fprintf(stderr, "Error: StarVLA does not support --noise-mode debug-sin; use Gaussian noise\n"); + return false; + } if (args.model_type == robotcpp::model_type::smolvla) { if (args.llm_path.empty() || args.mmproj_path.empty() || args.state_proj_path.empty() || args.action_expert_path.empty()) { std::fprintf(stderr, "Error: smolvla requires --llm --mmproj --state-proj --action-expert\n"); return false; } - } else if (args.vit_path.empty() || args.mmproj_path.empty() || args.llm_path.empty() || - args.tokenizer_path.empty() || args.state_path.empty() || args.action_decoder_path.empty()) { - std::fprintf(stderr, "Error: pi0 requires --vit --mmproj --llm --tokenizer --state-gguf --action-decoder\n"); + } else if (args.model_type == robotcpp::model_type::pi0) { + if (args.vit_path.empty() || args.mmproj_path.empty() || args.llm_path.empty() || args.tokenizer_path.empty() || + args.state_path.empty() || args.action_decoder_path.empty()) { + std::fprintf(stderr, + "Error: pi0 requires --vit --mmproj --llm --tokenizer --state-gguf --action-decoder\n"); + return false; + } + } else if (robotcpp::is_starvla_model_type(args.model_type)) { + if (args.llm_path.empty() || args.mmproj_path.empty() || args.policy_path.empty()) { + std::fprintf(stderr, "Error: %s requires --llm --mmproj --policy\n", + robotcpp::model_type_name(args.model_type)); + return false; + } + } else { + std::fprintf(stderr, "Error: unsupported model type '%s'\n", robotcpp::model_type_name(args.model_type)); return false; } return true; @@ -195,6 +239,7 @@ static robotcpp::model_args make_model_args(const server_args & args) { model_args.action_decoder_path = args.action_decoder_path; model_args.state_proj_path = args.state_proj_path; model_args.action_expert_path = args.action_expert_path; + model_args.policy_path = args.policy_path; model_args.n_batch = args.n_batch; model_args.n_ctx = args.n_ctx; model_args.noise_mode = args.noise_mode; diff --git a/robot_server/model_adapter.cpp b/robot_server/model_adapter.cpp index 31556ad..4b1a7b8 100644 --- a/robot_server/model_adapter.cpp +++ b/robot_server/model_adapter.cpp @@ -30,8 +30,9 @@ bool model_adapter::predict(const proto::predict_request & req, proto::predict_r image.stride_bytes = static_cast(src.stride_bytes); obs.images.push_back(image); } - obs.state = req.state; - obs.task = req.task; + obs.state = req.state; + obs.initial_noise = req.initial_noise; + obs.task = req.task; robotcpp::model_result result; if (!model_->predict(obs, result, error)) { diff --git a/robot_server/protocol.cpp b/robot_server/protocol.cpp index 8ea4c48..930a7db 100644 --- a/robot_server/protocol.cpp +++ b/robot_server/protocol.cpp @@ -1,5 +1,6 @@ #include "protocol.h" +#include #include #include @@ -122,6 +123,16 @@ static bool checked_u32_count(size_t n, const char * label, std::string & error) return true; } +static bool validate_f32_array(const std::vector & values, const char * label, std::string & error) { + for (float value : values) { + if (!std::isfinite(value)) { + error = std::string(label) + " contains a non-finite value"; + return false; + } + } + return true; +} + static bool validate_image_payload(const image_payload & image, const char * label, std::string & error) { if (image.image_format != image_raw_rgb_u8) { error = std::string(label) + " unsupported image format"; @@ -216,7 +227,13 @@ bool encode_predict_request(const predict_request & req, std::vector & return false; } if (!checked_u32_count(req.images.size(), "images", error) || - !checked_u32_count(req.state.size(), "state", error) || !checked_u32_count(req.task.size(), "task", error)) { + !checked_u32_count(req.state.size(), "state", error) || + !checked_u32_count(req.initial_noise.size(), "initial noise", error) || + !checked_u32_count(req.task.size(), "task", error)) { + return false; + } + if (!validate_f32_array(req.state, "state", error) || + !validate_f32_array(req.initial_noise, "initial noise", error)) { return false; } for (size_t i = 0; i < req.images.size(); ++i) { @@ -228,6 +245,7 @@ bool encode_predict_request(const predict_request & req, std::vector & put_u32(out, (uint32_t)req.images.size()); put_u32(out, (uint32_t)req.state.size()); + put_u32(out, (uint32_t)req.initial_noise.size()); put_u32(out, (uint32_t)req.task.size()); for (const image_payload & image : req.images) { put_u32(out, image.image_format); @@ -242,6 +260,9 @@ bool encode_predict_request(const predict_request & req, std::vector & for (float v : req.state) { put_f32(out, v); } + for (float v : req.initial_noise) { + put_f32(out, v); + } out.insert(out.end(), req.task.begin(), req.task.end()); for (const image_payload & image : req.images) { out.insert(out.end(), image.name.begin(), image.name.end()); @@ -255,9 +276,10 @@ bool decode_predict_request(const std::vector & payload, predict_reques reader r(payload.data(), payload.size()); uint32_t image_count = 0; uint32_t state_dim = 0; + uint32_t noise_dim = 0; uint32_t task_len = 0; - if (!r.u32(image_count) || !r.u32(state_dim) || !r.u32(task_len)) { + if (!r.u32(image_count) || !r.u32(state_dim) || !r.u32(noise_dim) || !r.u32(task_len)) { error = "short predict request"; return false; } @@ -265,6 +287,11 @@ bool decode_predict_request(const std::vector & payload, predict_reques error = "predict request requires at least one image"; return false; } + constexpr size_t image_metadata_size = 6 * sizeof(uint32_t) + sizeof(uint64_t); + if (image_count > r.remaining() / image_metadata_size) { + error = "image count exceeds predict request metadata"; + return false; + } req.images.assign(image_count, image_payload{}); std::vector name_lens(image_count, 0); @@ -286,6 +313,12 @@ bool decode_predict_request(const std::vector & payload, predict_reques } } + const uint64_t scalar_bytes = (static_cast(state_dim) + static_cast(noise_dim)) * sizeof(float); + if (scalar_bytes > r.remaining() || task_len > r.remaining() - scalar_bytes) { + error = "predict request fields exceed payload"; + return false; + } + req.state.assign(state_dim, 0.0f); for (uint32_t i = 0; i < state_dim; ++i) { if (!r.f32(req.state[i])) { @@ -293,6 +326,19 @@ bool decode_predict_request(const std::vector & payload, predict_reques return false; } } + if (!validate_f32_array(req.state, "state", error)) { + return false; + } + req.initial_noise.assign(noise_dim, 0.0f); + for (uint32_t i = 0; i < noise_dim; ++i) { + if (!r.f32(req.initial_noise[i])) { + error = "short initial noise array"; + return false; + } + } + if (!validate_f32_array(req.initial_noise, "initial noise", error)) { + return false; + } if (!r.string(req.task, task_len)) { error = "short task string"; return false; diff --git a/robot_server/protocol.h b/robot_server/protocol.h index 97ae387..a5f483e 100644 --- a/robot_server/protocol.h +++ b/robot_server/protocol.h @@ -9,7 +9,7 @@ namespace robot_server { namespace protocol { static constexpr uint32_t k_magic = 0x414c5653u; // "SVLA" in little-endian bytes. -static constexpr uint16_t k_version = 3; +static constexpr uint16_t k_version = 4; static constexpr uint16_t k_header_size = 32; static constexpr uint64_t k_default_max_payload = 256ull * 1024ull * 1024ull; @@ -63,6 +63,7 @@ struct metric { struct predict_request { std::vector images; std::vector state; + std::vector initial_noise; std::string task; }; diff --git a/robot_server/shell/launch_robot_server_linux_cuda.sh b/robot_server/shell/launch_robot_server_linux_cuda.sh index 990ed44..c6002b9 100755 --- a/robot_server/shell/launch_robot_server_linux_cuda.sh +++ b/robot_server/shell/launch_robot_server_linux_cuda.sh @@ -21,6 +21,10 @@ SKIP_BUILD="${SKIP_BUILD:-0}" CMAKE_BIN="${CMAKE_BIN:-cmake}" GGML_NATIVE="${GGML_NATIVE:-OFF}" GGML_OPENMP="${GGML_OPENMP:-OFF}" +ROBOT_CPP_BUILD_STARVLA="${ROBOT_CPP_BUILD_STARVLA:-OFF}" +if [ "${MODEL_TYPE}" = "starvla" ]; then + ROBOT_CPP_BUILD_STARVLA=ON +fi SERVER_BIN="${BUILD_DIR}/bin/model-server" @@ -32,7 +36,8 @@ if [ "${SKIP_BUILD}" != "1" ]; then -DGGML_OPENMP="${GGML_OPENMP}" \ -DGGML_CUDA=ON \ -DGGML_METAL=OFF \ - -DROBOT_CPP_BUILD_ROBOT_SERVER=ON + -DROBOT_CPP_BUILD_ROBOT_SERVER=ON \ + -DROBOT_CPP_BUILD_STARVLA="${ROBOT_CPP_BUILD_STARVLA}" echo "== build ==" "${CMAKE_BIN}" --build "${BUILD_DIR}" --target model-server -j8 @@ -70,6 +75,17 @@ case "${MODEL_TYPE}" in --action-decoder "${ACTION_DECODER_GGUF}" ) ;; + starvla) + LLM_GGUF="${LLM_GGUF:?LLM_GGUF must be set for StarVLA}" + MMPROJ_GGUF="${MMPROJ_GGUF:?MMPROJ_GGUF must be set for StarVLA}" + POLICY_GGUF="${POLICY_GGUF:?POLICY_GGUF must be set for StarVLA}" + MODEL_ARGS=( + --model-type starvla + --llm "${LLM_GGUF}" + --mmproj "${MMPROJ_GGUF}" + --policy "${POLICY_GGUF}" + ) + ;; *) echo "unsupported MODEL_TYPE=${MODEL_TYPE}" >&2 exit 1 diff --git a/robot_server/test/benchmark_latency.py b/robot_server/test/benchmark_latency.py index e000113..fa2ef72 100644 --- a/robot_server/test/benchmark_latency.py +++ b/robot_server/test/benchmark_latency.py @@ -26,8 +26,15 @@ def make_random_state(dim: int, seed: int) -> np.ndarray: return rng.uniform(-1.0, 1.0, size=(dim,)).astype(np.float32) -def make_random_observation(width: int, height: int, state_dim: int, prompt: str, image_names: list[str]) -> dict: - return { +def make_random_observation( + width: int, + height: int, + state_dim: int, + initial_noise_dim: int, + prompt: str, + image_names: list[str], +) -> dict: + observation = { "images": [ { "name": image_name, @@ -38,6 +45,9 @@ def make_random_observation(width: int, height: int, state_dim: int, prompt: str "state": make_random_state(state_dim, seed=1), "prompt": prompt, } + if initial_noise_dim: + observation["initial_noise"] = make_random_state(initial_noise_dim, seed=2) + return observation def ordered_columns(rows: list[dict[str, float]]) -> list[str]: @@ -141,6 +151,7 @@ def main() -> int: parser.add_argument("--height", type=int, default=224) parser.add_argument("--image-name", action="append") parser.add_argument("--state-dim", type=int, default=6) + parser.add_argument("--initial-noise-dim", type=int, default=0) parser.add_argument("--prompt", default=os.environ.get("SMOLVLA_PROMPT", "grab the block.")) parser.add_argument("--warmup", type=int, default=1) parser.add_argument("--loops", type=int, default=10) @@ -167,6 +178,7 @@ def main() -> int: width=args.width, height=args.height, state_dim=args.state_dim, + initial_noise_dim=args.initial_noise_dim, prompt=args.prompt, image_names=image_names, ) diff --git a/robot_server/test/test_server_latency.sh b/robot_server/test/test_server_latency.sh index 3073dff..de19516 100755 --- a/robot_server/test/test_server_latency.sh +++ b/robot_server/test/test_server_latency.sh @@ -7,9 +7,9 @@ set -e # bash robot_server/test/test_server_latency.sh # # Positional args: -# $1: model-type, e.g. smolvla / pi0 +# $1: model-type, e.g. smolvla / pi0 / starvla # $2: backend, e.g. mac-cpu / mac-metal / linux-cpu / linux-cuda -# $3: test-suite, e.g. smolvla-libero / smolvla-so101 / pi0-libero +# $3: test-suite, e.g. smolvla-libero / smolvla-so101 / pi0-libero / starvla-bridge ROBOT_CPP_ROOT="${ROBOT_CPP_ROOT:?ROBOT_CPP_ROOT must be set}" GGUF_DIR="${GGUF_DIR:?GGUF_DIR must be set}" MODEL_TYPE="${1:-${MODEL_TYPE:-smolvla}}" @@ -39,6 +39,17 @@ case "${MODEL_TYPE}" in ACTION_DECODER_GGUF="${ACTION_DECODER_GGUF:-${GGUF_DIR}/${MODEL_BASENAME}.action_decoder.gguf}" LLM_GGUF="${LLM_GGUF:-${GGUF_DIR}/${MODEL_BASENAME}.llm.gguf}" ;; + starvla) + mapfile -t LLM_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name 'qwen-*.gguf' -print) + mapfile -t MMPROJ_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name 'mmproj-*.gguf' -print) + mapfile -t POLICY_CANDIDATES < <(find "${GGUF_DIR}" -maxdepth 1 -type f -name '*policy*.gguf' -print) + [[ ${#LLM_CANDIDATES[@]} -eq 1 ]] || { echo "expected one Qwen GGUF in ${GGUF_DIR}" >&2; exit 1; } + [[ ${#MMPROJ_CANDIDATES[@]} -eq 1 ]] || { echo "expected one mmproj GGUF in ${GGUF_DIR}" >&2; exit 1; } + [[ ${#POLICY_CANDIDATES[@]} -eq 1 ]] || { echo "expected one policy GGUF in ${GGUF_DIR}" >&2; exit 1; } + LLM_GGUF="${LLM_GGUF:-${LLM_CANDIDATES[0]}}" + MMPROJ_GGUF="${MMPROJ_GGUF:-${MMPROJ_CANDIDATES[0]}}" + POLICY_GGUF="${POLICY_GGUF:-${POLICY_CANDIDATES[0]}}" + ;; *) echo "unsupported MODEL_TYPE=${MODEL_TYPE}" >&2 exit 1 @@ -99,6 +110,12 @@ case "${TEST_SUITE}" in IMAGE_HEIGHT="${IMAGE_HEIGHT:-256}" STATE_DIM="${STATE_DIM:-8}" ;; + starvla-bridge) + IMAGE_NAMES="${IMAGE_NAMES:-${IMAGE_NAME:-image_0}}" + IMAGE_WIDTH="${IMAGE_WIDTH:-224}" + IMAGE_HEIGHT="${IMAGE_HEIGHT:-224}" + STATE_DIM="${STATE_DIM:-0}" + ;; *) echo "unsupported TEST_SUITE=${TEST_SUITE}" >&2 exit 1 @@ -108,6 +125,7 @@ WARMUP="${WARMUP:-5}" LOOPS="${LOOPS:-100}" SERVER_WAIT_S="${SERVER_WAIT_S:-120}" DTYPE="${DTYPE:-f32}" +NOISE_SEED="${NOISE_SEED:--1}" PYTHON="${PYTHON:-python3}" # ==================================== @@ -155,12 +173,13 @@ run_latency_case() { TOKENIZER_GGUF="${TOKENIZER_GGUF:-}" \ STATE_GGUF="${STATE_GGUF:-}" \ ACTION_DECODER_GGUF="${ACTION_DECODER_GGUF:-}" \ + POLICY_GGUF="${POLICY_GGUF:-}" \ HOST="${HOST}" \ PORT="${PORT}" \ THREADS="${threads}" \ TASK="${PROMPT}" \ NOISE_MODE="gaussian" \ - NOISE_SEED="-1" \ + NOISE_SEED="${NOISE_SEED}" \ bash "${LAUNCH_SHELL}" "${MODEL_TYPE}" >"${server_log}" 2>&1 & SERVER_PID=$! diff --git a/src/model-cli.cpp b/src/model-cli.cpp index bc82d8e..8dc8441 100644 --- a/src/model-cli.cpp +++ b/src/model-cli.cpp @@ -1,7 +1,9 @@ // model-cli.cpp — common robotcpp::Model CLI frontend #include "models/model.h" +#include "models/argument_parse.h" #include "models/smolvla/smolvla_engine.h" +#include "llama.h" #include "stb_image.h" #include @@ -26,16 +28,11 @@ struct loaded_image { int stride_bytes = 0; }; -bool parse_model_type(const std::string & value, robotcpp::model_type & out) { - if (value == "smolvla") { - out = robotcpp::model_type::smolvla; - return true; - } - if (value == "pi0") { - out = robotcpp::model_type::pi0; - return true; +void quiet_llama_log_callback(ggml_log_level level, const char * text, void * user_data) { + (void)user_data; + if (level == GGML_LOG_LEVEL_ERROR) { + std::fputs(text, stderr); } - return false; } bool parse_noise_mode(const std::string & value, int & out_mode) { @@ -54,13 +51,14 @@ void print_usage(const char * prog) { std::fprintf(stderr, "\nModel CLI - robotcpp::Model frontend\n\n"); std::fprintf(stderr, "Usage:\n"); std::fprintf(stderr, " %s --model-type smolvla [options]\n", prog); - std::fprintf(stderr, " %s --model-type pi0 [options]\n\n", prog); + std::fprintf(stderr, " %s --model-type pi0 [options]\n", prog); + std::fprintf(stderr, " %s --model-type starvla --llm --mmproj --policy [options]\n\n", prog); std::fprintf(stderr, "Common options:\n"); - std::fprintf(stderr, " --model-type Model type (default: smolvla)\n"); + std::fprintf(stderr, " --model-type smolvla|pi0|starvla\n" + " (default: smolvla)\n"); std::fprintf(stderr, " --image Input image (repeatable; order matches --image-name)\n"); - std::fprintf( - stderr, - " --image-name Observation image name (repeatable; default: image for single-image input)\n"); + std::fprintf(stderr, + " --image-name Observation image name (default: image_0 for StarVLA, image otherwise)\n"); std::fprintf(stderr, " --state Proprio/state values (comma-separated)\n"); std::fprintf(stderr, " --task Task instruction (default: \"grab the block.\")\n"); std::fprintf(stderr, " --threads Number of threads (default: auto)\n"); @@ -82,6 +80,13 @@ void print_usage(const char * prog) { std::fprintf(stderr, " --tokenizer Tokenizer GGUF path\n"); std::fprintf(stderr, " --state-gguf State projector GGUF path\n"); std::fprintf(stderr, " --action-decoder Action decoder GGUF path\n"); + std::fprintf(stderr, "\nStarVLA options:\n"); + std::fprintf(stderr, " --policy StarVLA policy GGUF path (required)\n"); + std::fprintf(stderr, " --llm Qwen text GGUF path (required)\n"); + std::fprintf(stderr, " --mmproj Qwen vision GGUF path (required)\n"); + std::fprintf(stderr, " --n-batch Qwen batch size (default: 512)\n"); + std::fprintf(stderr, " --n-ctx Qwen context size (default: 2048)\n"); + std::fprintf(stderr, " --noise-seed GR00T/PI/PI_v3 noise seed, <0 means auto (default: -1)\n"); } bool parse_state(const char * csv, std::vector & out) { @@ -128,7 +133,7 @@ int main(int argc, char ** argv) { std::vector image_paths; std::vector image_names; std::string state_csv; - std::string task; + std::string task = "grab the block."; for (int i = 1; i < argc; i++) { std::string arg = argv[i]; @@ -139,10 +144,12 @@ int main(int argc, char ** argv) { } else if (arg == "-v" || arg == "--verbose") { args.verbosity++; } else if (arg == "--model-type" && i + 1 < argc) { - if (!parse_model_type(argv[++i], args.type)) { + if (!robotcpp::parse_model_type(argv[++i], args.type)) { std::fprintf(stderr, "Error: unsupported model type '%s'\n", argv[i]); return 1; } + } else if (arg == "--policy" && i + 1 < argc) { + args.policy_path = argv[++i]; } else if (arg == "--llm" && i + 1 < argc) { args.llm_path = argv[++i]; } else if (arg == "--mmproj" && i + 1 < argc) { @@ -168,18 +175,34 @@ int main(int argc, char ** argv) { } else if (arg == "--task" && i + 1 < argc) { task = argv[++i]; } else if (arg == "--threads" && i + 1 < argc) { - args.threads = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.threads)) { + std::fprintf(stderr, "Error: invalid --threads value '%s'\n", value); + return 1; + } } else if (arg == "--n-batch" && i + 1 < argc) { - args.n_batch = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_batch)) { + std::fprintf(stderr, "Error: invalid --n-batch value '%s'\n", value); + return 1; + } } else if (arg == "--n-ctx" && i + 1 < argc) { - args.n_ctx = std::atoi(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.n_ctx)) { + std::fprintf(stderr, "Error: invalid --n-ctx value '%s'\n", value); + return 1; + } } else if (arg == "--noise-mode" && i + 1 < argc) { if (!parse_noise_mode(argv[++i], args.noise_mode)) { std::fprintf(stderr, "Error: invalid noise mode '%s'\n", argv[i]); return 1; } } else if (arg == "--noise-seed" && i + 1 < argc) { - args.noise_seed = std::atoll(argv[++i]); + const char * value = argv[++i]; + if (!robotcpp::parse_integer_argument(value, args.noise_seed)) { + std::fprintf(stderr, "Error: invalid --noise-seed value '%s'\n", value); + return 1; + } } else { std::fprintf(stderr, "Error: unknown argument '%s'\n", arg.c_str()); print_usage(argv[0]); @@ -187,23 +210,46 @@ int main(int argc, char ** argv) { } } + if (args.threads < 0 || args.n_batch <= 0 || args.n_ctx <= 0) { + std::fprintf(stderr, "Error: --threads must be non-negative and --n-batch/--n-ctx must be positive\n"); + return 1; + } + if (robotcpp::is_starvla_model_type(args.type)) { + if (args.llm_path.empty() || args.mmproj_path.empty() || args.policy_path.empty()) { + std::fprintf(stderr, "Error: %s requires --llm --mmproj --policy\n", robotcpp::model_type_name(args.type)); + return 1; + } + if (args.noise_mode != SMOLVLA_NOISE_MODE_GAUSSIAN) { + std::fprintf(stderr, "Error: StarVLA does not support --noise-mode debug-sin; use Gaussian noise\n"); + return 1; + } + } + if (image_paths.empty()) { std::fprintf(stderr, "Error: --image is required\n"); print_usage(argv[0]); return 1; } + if (robotcpp::is_starvla_model_type(args.type) && image_paths.size() != 1) { + std::fprintf(stderr, "Error: %s requires exactly one --image\n", robotcpp::model_type_name(args.type)); + return 1; + } if (image_names.empty()) { if (image_paths.size() != 1) { std::fprintf(stderr, "Error: multiple --image inputs require one --image-name per image\n"); return 1; } - image_names.push_back("image"); + image_names.push_back(robotcpp::is_starvla_model_type(args.type) ? "image_0" : "image"); } if (image_names.size() != image_paths.size()) { std::fprintf(stderr, "Error: --image count (%zu) must match --image-name count (%zu)\n", image_paths.size(), image_names.size()); return 1; } + if (robotcpp::is_starvla_model_type(args.type) && image_names[0] != "image_0") { + std::fprintf(stderr, "Error: %s image must be named 'image_0'\n", robotcpp::model_type_name(args.type)); + return 1; + } std::vector state_vec; if (!parse_state(state_csv.c_str(), state_vec)) { @@ -218,6 +264,8 @@ int main(int argc, char ** argv) { } } + llama_log_set(args.verbosity > 0 ? nullptr : quiet_llama_log_callback, nullptr); + const auto init_start = std::chrono::high_resolution_clock::now(); std::string error; std::unique_ptr model; diff --git a/src/models/argument_parse.h b/src/models/argument_parse.h new file mode 100644 index 0000000..0ad5b60 --- /dev/null +++ b/src/models/argument_parse.h @@ -0,0 +1,27 @@ +#pragma once + +#include +#include +#include +#include + +namespace robotcpp { + +template bool parse_integer_argument(const char * value, Integer & output) { + static_assert(std::is_integral::value && !std::is_same::value, + "Integer must be a non-bool integral type"); + if (value == nullptr || value[0] == '\0') { + return false; + } + + Integer parsed = 0; + const char * end = value + std::strlen(value); + const std::from_chars_result result = std::from_chars(value, end, parsed, 10); + if (result.ec != std::errc{} || result.ptr != end) { + return false; + } + output = parsed; + return true; +} + +} // namespace robotcpp diff --git a/src/models/ggml_backend.cpp b/src/models/ggml_backend.cpp index 062d895..65ef5ed 100644 --- a/src/models/ggml_backend.cpp +++ b/src/models/ggml_backend.cpp @@ -6,7 +6,7 @@ #include #include -static const char * backend_mode_name(backend_mode mode) { +const char * backend_mode_name(backend_mode mode) { switch (mode) { case backend_mode::cuda: return "cuda"; diff --git a/src/models/ggml_backend.h b/src/models/ggml_backend.h index 08b735d..a044f6f 100644 --- a/src/models/ggml_backend.h +++ b/src/models/ggml_backend.h @@ -11,6 +11,8 @@ enum class backend_mode { metal, }; +const char * backend_mode_name(backend_mode mode); + struct backend_buft_policy { ggml_backend_buffer_type_t model_buft = nullptr; ggml_backend_buffer_type_t runtime_buft = nullptr; diff --git a/src/models/model.h b/src/models/model.h index 587957c..d89b9d3 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -10,8 +10,13 @@ namespace robotcpp { enum class model_type { smolvla, pi0, + starvla, }; +const char * model_type_name(model_type type); +bool parse_model_type(const std::string & value, model_type & out); +bool is_starvla_model_type(model_type type); + struct model_image { std::string name; const uint8_t * data = nullptr; @@ -24,6 +29,7 @@ struct model_image { struct observation { std::vector images; std::vector state; + std::vector initial_noise; std::string task; }; @@ -59,6 +65,9 @@ struct model_args { std::string tokenizer_path; std::string state_path; std::string action_decoder_path; + + // starvla + std::string policy_path; }; class Model { diff --git a/src/models/model_factory.cpp b/src/models/model_factory.cpp index d6b1aa3..3f887ad 100644 --- a/src/models/model_factory.cpp +++ b/src/models/model_factory.cpp @@ -2,6 +2,9 @@ #include "models/pi0/pi0_model.h" #include "models/smolvla/smolvla_model.h" +#ifdef ROBOT_CPP_BUILD_STARVLA +#include "models/starvla/starvla_model.h" +#endif namespace robotcpp { @@ -13,8 +16,16 @@ bool make_model(const model_args & args, std::unique_ptr & out, std::stri if (args.type == model_type::pi0) { return make_pi0_model(args, out, error); } + if (args.type == model_type::starvla) { +#ifdef ROBOT_CPP_BUILD_STARVLA + return make_starvla_model(args, out, error); +#else + error = "StarVLA support was not built; configure with -DROBOT_CPP_BUILD_STARVLA=ON"; + return false; +#endif + } - error = "unsupported model type"; + error = std::string("unsupported model type: ") + model_type_name(args.type); return false; } diff --git a/src/models/model_type.cpp b/src/models/model_type.cpp new file mode 100644 index 0000000..6a4d3d5 --- /dev/null +++ b/src/models/model_type.cpp @@ -0,0 +1,49 @@ +#include "models/model.h" + +#include + +namespace robotcpp { +namespace { + +struct model_type_entry { + model_type type; + const char * name; +}; + +constexpr std::array MODEL_TYPES = {{ + {model_type::smolvla, "smolvla"}, + {model_type::pi0, "pi0"}, + {model_type::starvla, "starvla"}, +}}; + +const model_type_entry * find_entry(model_type type) { + for (const model_type_entry & entry : MODEL_TYPES) { + if (entry.type == type) { + return &entry; + } + } + return nullptr; +} + +} // namespace + +const char * model_type_name(model_type type) { + const model_type_entry * entry = find_entry(type); + return entry ? entry->name : "unknown"; +} + +bool parse_model_type(const std::string & value, model_type & out) { + for (const model_type_entry & entry : MODEL_TYPES) { + if (value == entry.name) { + out = entry.type; + return true; + } + } + return false; +} + +bool is_starvla_model_type(model_type type) { + return type == model_type::starvla; +} + +} // namespace robotcpp diff --git a/src/models/pi0/pi0_model.cpp b/src/models/pi0/pi0_model.cpp index 5d62866..4d46a0d 100644 --- a/src/models/pi0/pi0_model.cpp +++ b/src/models/pi0/pi0_model.cpp @@ -80,6 +80,10 @@ bool Pi0Model::predict(const observation & obs, model_result & out, std::string error = "Pi0 model is not initialized"; return false; } + if (!obs.initial_noise.empty()) { + error = "Pi0 does not accept explicit initial noise"; + return false; + } if (obs.images.empty()) { error = "Pi0 requires at least one image"; return false; diff --git a/src/models/smolvla/smolvla_model.cpp b/src/models/smolvla/smolvla_model.cpp index a718fd2..db6567c 100644 --- a/src/models/smolvla/smolvla_model.cpp +++ b/src/models/smolvla/smolvla_model.cpp @@ -72,6 +72,10 @@ bool SmolVLAModel::predict(const observation & obs, model_result & out, std::str error = "SmolVLA model is not initialized"; return false; } + if (!obs.initial_noise.empty()) { + error = "SmolVLA does not accept explicit initial noise"; + return false; + } if (obs.images.empty()) { error = "SmolVLA requires at least one image"; return false; diff --git a/src/models/starvla/fast_codec.cpp b/src/models/starvla/fast_codec.cpp new file mode 100644 index 0000000..9ce67fd --- /dev/null +++ b/src/models/starvla/fast_codec.cpp @@ -0,0 +1,537 @@ +#include "models/starvla/fast_codec.h" + +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +constexpr size_t kMaximumVocabSize = 65536U; +constexpr size_t kMaximumTimeHorizon = 1024U; +constexpr size_t kMaximumActionDim = 1024U; +constexpr size_t kMaximumBatchSize = 1024U; +constexpr size_t kMaximumTokenSequence = 4096U; +constexpr size_t kMaximumGeneratedSequence = 2048U; +constexpr size_t kMaximumDecodedBytes = 1024U * 1024U; +constexpr size_t kMaximumOutputScalars = 16U * 1024U * 1024U; +constexpr uint64_t kMaximumIdctMultiplyAdds = 64ULL * 1024ULL * 1024ULL; +constexpr double kPi = 3.141592653589793238462643383279502884; + +bool decode_utf8_strict(const std::string & input, std::vector & output, std::string & error) { + output.clear(); + for (size_t i = 0; i < input.size();) { + const uint8_t first = static_cast(input[i]); + uint32_t value = 0; + size_t length = 0; + if (first <= 0x7fU) { + value = first; + length = 1; + } else if (first >= 0xc2U && first <= 0xdfU) { + value = first & 0x1fU; + length = 2; + } else if (first >= 0xe0U && first <= 0xefU) { + value = first & 0x0fU; + length = 3; + } else if (first >= 0xf0U && first <= 0xf4U) { + value = first & 0x07U; + length = 4; + } else { + error = "StarVLA FAST tokenizer vocabulary contains invalid UTF-8"; + return false; + } + if (i + length > input.size()) { + error = "StarVLA FAST tokenizer vocabulary contains truncated UTF-8"; + return false; + } + for (size_t j = 1; j < length; ++j) { + const uint8_t continuation = static_cast(input[i + j]); + if ((continuation & 0xc0U) != 0x80U) { + error = "StarVLA FAST tokenizer vocabulary contains invalid UTF-8 continuation"; + return false; + } + value = (value << 6U) | (continuation & 0x3fU); + } + const bool overlong = + (length == 2 && value < 0x80U) || (length == 3 && value < 0x800U) || (length == 4 && value < 0x10000U); + if (overlong || value > 0x10ffffU || (value >= 0xd800U && value <= 0xdfffU)) { + error = "StarVLA FAST tokenizer vocabulary contains a non-scalar UTF-8 value"; + return false; + } + output.push_back(value); + i += length; + } + return true; +} + +std::unordered_map byte_level_inverse_alphabet() { + std::unordered_map result; + std::unordered_set direct; + for (int value = 0x21; value <= 0x7e; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + for (int value = 0xa1; value <= 0xac; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + for (int value = 0xae; value <= 0xff; ++value) { + direct.insert(value); + result.emplace(static_cast(value), static_cast(value)); + } + uint32_t extra = 0; + for (int value = 0; value <= 0xff; ++value) { + if (direct.count(value) == 0) { + result.emplace(256U + extra, static_cast(value)); + ++extra; + } + } + return result; +} + +bool compile_token_bytes(const std::vector & vocab_by_id, std::vector> & token_bytes, + std::string & error) { + const auto inverse_alphabet = byte_level_inverse_alphabet(); + token_bytes.clear(); + token_bytes.reserve(vocab_by_id.size()); + for (size_t token_id = 0; token_id < vocab_by_id.size(); ++token_id) { + if (vocab_by_id[token_id].empty()) { + error = "StarVLA FAST tokenizer has an empty vocabulary piece at ID " + std::to_string(token_id); + return false; + } + std::vector piece_codepoints; + if (!decode_utf8_strict(vocab_by_id[token_id], piece_codepoints, error)) { + error += " at token ID " + std::to_string(token_id); + return false; + } + std::vector bytes; + bytes.reserve(piece_codepoints.size()); + for (uint32_t codepoint : piece_codepoints) { + const auto found = inverse_alphabet.find(codepoint); + if (found == inverse_alphabet.end()) { + error = "StarVLA FAST tokenizer piece contains a code point outside the ByteLevel " + "alphabet at ID " + + std::to_string(token_id); + return false; + } + bytes.push_back(found->second); + } + token_bytes.push_back(std::move(bytes)); + } + return true; +} + +void decode_utf8_lossy(const std::vector & input, std::vector & output) { + output.clear(); + for (size_t i = 0; i < input.size();) { + const uint8_t first = input[i]; + if (first <= 0x7fU) { + output.push_back(first); + ++i; + continue; + } + + size_t length = 0; + uint32_t value = 0; + if (first >= 0xc2U && first <= 0xdfU) { + length = 2; + value = first & 0x1fU; + } else if (first >= 0xe0U && first <= 0xefU) { + length = 3; + value = first & 0x0fU; + } else if (first >= 0xf0U && first <= 0xf4U) { + length = 4; + value = first & 0x07U; + } else { + output.push_back(0xfffdU); + ++i; + continue; + } + + if (i + 1 >= input.size()) { + output.push_back(0xfffdU); + break; + } + const uint8_t second = input[i + 1]; + const bool second_is_continuation = (second & 0xc0U) == 0x80U; + const bool second_in_scalar_range = !(first == 0xe0U && second < 0xa0U) && + !(first == 0xedU && second > 0x9fU) && + !(first == 0xf0U && second < 0x90U) && !(first == 0xf4U && second > 0x8fU); + if (!second_is_continuation || !second_in_scalar_range) { + output.push_back(0xfffdU); + ++i; + continue; + } + value = (value << 6U) | (second & 0x3fU); + + bool invalid = false; + size_t consumed_prefix = 2; + for (size_t j = 2; j < length; ++j) { + if (i + j >= input.size()) { + output.push_back(0xfffdU); + i = input.size(); + invalid = true; + break; + } + const uint8_t continuation = input[i + j]; + if ((continuation & 0xc0U) != 0x80U) { + output.push_back(0xfffdU); + i += consumed_prefix; + invalid = true; + break; + } + value = (value << 6U) | (continuation & 0x3fU); + ++consumed_prefix; + } + if (invalid) { + continue; + } + output.push_back(value); + i += length; + } +} + +bool checked_action_count(const FastCodecConfig & config, size_t batch_size, size_t & per_sample, size_t & total, + std::string & error) { + if (config.vocab_size > kMaximumVocabSize || config.time_horizon > kMaximumTimeHorizon || + config.action_dim > kMaximumActionDim) { + error = "StarVLA FAST codec dimensions exceed the runtime safety limits"; + return false; + } + if (batch_size == 0 || batch_size > kMaximumBatchSize) { + error = "StarVLA FAST batch size exceeds the runtime safety limit"; + return false; + } + if (config.time_horizon > std::numeric_limits::max() / config.action_dim) { + error = "StarVLA FAST action shape overflows size_t"; + return false; + } + per_sample = config.time_horizon * config.action_dim; + if (batch_size > std::numeric_limits::max() / per_sample) { + error = "StarVLA FAST batch shape overflows size_t"; + return false; + } + total = batch_size * per_sample; + if (total > kMaximumOutputScalars) { + error = "StarVLA FAST output tensor exceeds the runtime scalar limit"; + return false; + } + const uint64_t horizon = static_cast(config.time_horizon); + const uint64_t action_dim = static_cast(config.action_dim); + const uint64_t batch = static_cast(batch_size); + if (horizon > kMaximumIdctMultiplyAdds / horizon) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + uint64_t multiply_adds = horizon * horizon; + if (action_dim > kMaximumIdctMultiplyAdds / multiply_adds) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + multiply_adds *= action_dim; + if (batch > kMaximumIdctMultiplyAdds / multiply_adds) { + error = "StarVLA FAST inverse DCT exceeds the runtime work limit"; + return false; + } + return true; +} + +} // namespace + +FastCodec::FastCodec(FastCodecConfig config, std::vector> token_bytes, + std::vector fast_to_vlm_id) + : config_(config), token_bytes_(std::move(token_bytes)), fast_to_vlm_id_(std::move(fast_to_vlm_id)) { + vlm_to_fast_id_.reserve(fast_to_vlm_id_.size()); + for (size_t fast_id = 0; fast_id < fast_to_vlm_id_.size(); ++fast_id) { + vlm_to_fast_id_.emplace_back(fast_to_vlm_id_[fast_id], static_cast(fast_id)); + } + std::sort(vlm_to_fast_id_.begin(), vlm_to_fast_id_.end()); +} + +std::unique_ptr FastCodec::create(FastCodecConfig config, std::vector vocab_by_id, + std::vector fast_to_vlm_id, std::string & error) { + error.clear(); + if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || config.time_horizon == 0 || + config.action_dim == 0) { + error = "StarVLA FAST codec dimensions and scale must be non-zero and finite"; + return nullptr; + } + if (config.vocab_size > static_cast(std::numeric_limits::max())) { + error = "StarVLA FAST vocabulary exceeds the int32 token-ID range"; + return nullptr; + } + if (vocab_by_id.size() != config.vocab_size || fast_to_vlm_id.size() != config.vocab_size) { + error = "StarVLA FAST codec vocabulary or action-token map has the wrong size"; + return nullptr; + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config, 1, per_sample, total, error)) { + return nullptr; + } + std::unordered_set unique_vlm_ids; + for (int32_t vlm_id : fast_to_vlm_id) { + if (vlm_id < 0 || !unique_vlm_ids.insert(vlm_id).second) { + error = "StarVLA FAST action-token VLM IDs must be unique and non-negative"; + return nullptr; + } + } + std::vector> token_bytes; + if (!compile_token_bytes(vocab_by_id, token_bytes, error)) { + return nullptr; + } + return std::unique_ptr(new FastCodec(config, std::move(token_bytes), std::move(fast_to_vlm_id))); +} + +std::unique_ptr FastCodec::create_compiled(FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, + std::vector fast_to_vlm_id, std::string & error) { + error.clear(); + if (!std::isfinite(config.scale) || config.scale == 0.0 || config.vocab_size == 0 || config.time_horizon == 0 || + config.action_dim == 0 || config.vocab_size > static_cast(std::numeric_limits::max())) { + error = "StarVLA FAST compiled codec dimensions and scale are invalid"; + return nullptr; + } + if (config.vocab_size == std::numeric_limits::max() || token_offsets.size() != config.vocab_size + 1U || + fast_to_vlm_id.size() != config.vocab_size || token_offsets.empty() || token_offsets.front() != 0 || + token_offsets.back() < 0 || static_cast(token_offsets.back()) != token_bytes.size()) { + error = "StarVLA FAST compiled codec tensor shapes are incompatible"; + return nullptr; + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config, 1, per_sample, total, error)) { + return nullptr; + } + + std::unordered_set unique_vlm_ids; + for (int32_t vlm_id : fast_to_vlm_id) { + if (vlm_id < 0 || !unique_vlm_ids.insert(vlm_id).second) { + error = "StarVLA FAST compiled action-token IDs must be unique and non-negative"; + return nullptr; + } + } + + std::vector> pieces; + pieces.reserve(config.vocab_size); + for (size_t index = 0; index < config.vocab_size; ++index) { + const int32_t begin = token_offsets[index]; + const int32_t end = token_offsets[index + 1U]; + if (begin < 0 || end <= begin || static_cast(end) > token_bytes.size()) { + error = "StarVLA FAST compiled codec offsets are not strictly increasing"; + return nullptr; + } + pieces.emplace_back(token_bytes.begin() + begin, token_bytes.begin() + end); + } + return std::unique_ptr(new FastCodec(config, std::move(pieces), std::move(fast_to_vlm_id))); +} + +const FastCodecConfig & FastCodec::config() const { + return config_; +} + +const std::vector & FastCodec::fast_to_vlm_ids() const { + return fast_to_vlm_id_; +} + +bool FastCodec::map_fast_to_vlm(const std::vector & fast_ids, std::vector & vlm_ids, + std::string & error) const { + vlm_ids.clear(); + error.clear(); + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + vlm_ids.reserve(fast_ids.size()); + for (int32_t fast_id : fast_ids) { + if (fast_id < 0 || static_cast(fast_id) >= fast_to_vlm_id_.size()) { + error = "StarVLA FAST token ID is outside the codec vocabulary"; + vlm_ids.clear(); + return false; + } + vlm_ids.push_back(fast_to_vlm_id_[static_cast(fast_id)]); + } + return true; +} + +bool FastCodec::map_vlm_to_fast(const std::vector & vlm_ids, std::vector & fast_ids, + std::string & error) const { + fast_ids.clear(); + error.clear(); + if (vlm_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST action-token sequence exceeds the runtime length limit"; + return false; + } + fast_ids.reserve(vlm_ids.size()); + for (int32_t vlm_id : vlm_ids) { + const auto found = std::lower_bound( + vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, + [](const std::pair & entry, int32_t value) { return entry.first < value; }); + if (found == vlm_to_fast_id_.end() || found->first != vlm_id) { + error = "Qwen token ID is not present in the StarVLA FAST action-token map"; + fast_ids.clear(); + return false; + } + fast_ids.push_back(found->second); + } + return true; +} + +bool FastCodec::extract_fast_tokens(const std::vector & generated_ids, std::vector & fast_ids, + std::string & error) const { + fast_ids.clear(); + error.clear(); + if (generated_ids.size() > kMaximumGeneratedSequence) { + error = "Qwen generated sequence exceeds the StarVLA FAST runtime length limit"; + return false; + } + for (int32_t vlm_id : generated_ids) { + const auto found = std::lower_bound( + vlm_to_fast_id_.begin(), vlm_to_fast_id_.end(), vlm_id, + [](const std::pair & entry, int32_t value) { return entry.first < value; }); + if (found != vlm_to_fast_id_.end() && found->first == vlm_id) { + fast_ids.push_back(found->second); + } + } + return true; +} + +bool FastCodec::byte_level_decode(const std::vector & fast_ids, std::vector & codepoints, + std::string & error) const { + codepoints.clear(); + error.clear(); + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + size_t byte_count = 0; + for (int32_t fast_id : fast_ids) { + if (fast_id < 0 || static_cast(fast_id) >= token_bytes_.size()) { + error = "StarVLA FAST token ID is outside the ByteLevel BPE vocabulary"; + return false; + } + const size_t piece_size = token_bytes_[static_cast(fast_id)].size(); + if (byte_count > std::numeric_limits::max() - piece_size) { + error = "StarVLA FAST ByteLevel output size overflows size_t"; + return false; + } + byte_count += piece_size; + if (byte_count > kMaximumDecodedBytes) { + error = "StarVLA FAST ByteLevel decode exceeds the runtime byte limit"; + return false; + } + } + std::vector bytes; + bytes.reserve(byte_count); + for (int32_t fast_id : fast_ids) { + const auto & piece = token_bytes_[static_cast(fast_id)]; + bytes.insert(bytes.end(), piece.begin(), piece.end()); + } + decode_utf8_lossy(bytes, codepoints); + return true; +} + +bool FastCodec::decode_fast_tokens(const std::vector> & batch_fast_ids, FastDecodeResult & result, + std::string & error) const { + result = {}; + error.clear(); + if (batch_fast_ids.empty()) { + error = "StarVLA FAST decode batch must contain at least one sequence"; + return false; + } + if (batch_fast_ids.size() > kMaximumBatchSize) { + error = "StarVLA FAST decode batch exceeds the runtime size limit"; + return false; + } + for (const auto & fast_ids : batch_fast_ids) { + if (fast_ids.size() > kMaximumTokenSequence) { + error = "StarVLA FAST token sequence exceeds the runtime length limit"; + return false; + } + } + size_t per_sample = 0; + size_t total = 0; + if (!checked_action_count(config_, batch_fast_ids.size(), per_sample, total, error)) { + return false; + } + + result.batch_size = batch_fast_ids.size(); + result.time_horizon = config_.time_horizon; + result.action_dim = config_.action_dim; + result.actions.assign(total, 0.0); + + const double dc_scale = 1.0 / std::sqrt(static_cast(config_.time_horizon)); + const double ac_scale = std::sqrt(2.0 / static_cast(config_.time_horizon)); + for (size_t batch = 0; batch < batch_fast_ids.size(); ++batch) { + std::vector codepoints; + std::string sequence_error; + if (!byte_level_decode(batch_fast_ids[batch], codepoints, sequence_error) || codepoints.size() != per_sample) { + error = "StarVLA FAST sequence " + std::to_string(batch) + ": " + + (sequence_error.empty() ? "decoded DCT coefficient shape mismatch" : sequence_error); + result = {}; + return false; + } + + for (size_t action = 0; action < config_.action_dim; ++action) { + const double dc = (static_cast(codepoints[action]) + config_.min_token) / config_.scale; + for (size_t time = 0; time < config_.time_horizon; ++time) { + double value = dc_scale * dc; + for (size_t frequency = 1; frequency < config_.time_horizon; ++frequency) { + const size_t coefficient_index = frequency * config_.action_dim + action; + const double coefficient = + (static_cast(codepoints[coefficient_index]) + config_.min_token) / config_.scale; + const double angle = kPi * static_cast(frequency) * static_cast(2U * time + 1U) / + (2.0 * static_cast(config_.time_horizon)); + value += ac_scale * coefficient * std::cos(angle); + } + result.actions[batch * per_sample + time * config_.action_dim + action] = value; + } + } + } + return true; +} + +bool FastCodec::decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, + FastDecodeResult & result, std::string & error) const { + if (batch_vlm_ids.empty() || batch_vlm_ids.size() > kMaximumBatchSize) { + result = {}; + error = "StarVLA FAST action-token batch is empty or exceeds the runtime size limit"; + return false; + } + std::vector> batch_fast_ids; + batch_fast_ids.reserve(batch_vlm_ids.size()); + for (const auto & vlm_ids : batch_vlm_ids) { + std::vector fast_ids; + if (!map_vlm_to_fast(vlm_ids, fast_ids, error)) { + result = {}; + return false; + } + batch_fast_ids.push_back(std::move(fast_ids)); + } + return decode_fast_tokens(batch_fast_ids, result, error); +} + +bool FastCodec::decode_generated_tokens(const std::vector> & batch_generated_ids, + FastDecodeResult & result, std::string & error) const { + if (batch_generated_ids.empty() || batch_generated_ids.size() > kMaximumBatchSize) { + result = {}; + error = "StarVLA FAST generated-token batch is empty or exceeds the runtime size limit"; + return false; + } + std::vector> batch_fast_ids; + batch_fast_ids.reserve(batch_generated_ids.size()); + for (const auto & generated_ids : batch_generated_ids) { + std::vector fast_ids; + if (!extract_fast_tokens(generated_ids, fast_ids, error)) { + result = {}; + return false; + } + batch_fast_ids.push_back(std::move(fast_ids)); + } + return decode_fast_tokens(batch_fast_ids, result, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_codec.h b/src/models/starvla/fast_codec.h new file mode 100644 index 0000000..4c3742b --- /dev/null +++ b/src/models/starvla/fast_codec.h @@ -0,0 +1,79 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct FastCodecConfig { + double scale = 0.0; + int32_t min_token = 0; + size_t vocab_size = 0; + size_t time_horizon = 0; + size_t action_dim = 0; +}; + +struct FastDecodeResult { + size_t batch_size = 0; + size_t time_horizon = 0; + size_t action_dim = 0; + std::vector actions; +}; + +class FastCodec { + public: + static std::unique_ptr create(FastCodecConfig config, std::vector vocab_by_id, + std::vector fast_to_vlm_id, std::string & error); + + // Constructs directly from the converter-compiled ByteLevel pieces stored + // in policy GGUF. offsets has vocab_size + 1 entries and indexes the flat + // byte buffer; no external tokenizer JSON is consulted. + static std::unique_ptr create_compiled(FastCodecConfig config, std::vector token_offsets, + std::vector token_bytes, + std::vector fast_to_vlm_id, std::string & error); + + const FastCodecConfig & config() const; + const std::vector & fast_to_vlm_ids() const; + + bool map_fast_to_vlm(const std::vector & fast_ids, std::vector & vlm_ids, + std::string & error) const; + bool map_vlm_to_fast(const std::vector & vlm_ids, std::vector & fast_ids, + std::string & error) const; + + // Extracts every mapped action token from a generated Qwen sequence in order. + // EOS stopping remains the generator's responsibility; ordinary EOS/pad/text + // IDs in the returned sequence are ignored and do not terminate this scan. + bool extract_fast_tokens(const std::vector & generated_ids, std::vector & fast_ids, + std::string & error) const; + + bool decode_fast_tokens(const std::vector> & batch_fast_ids, FastDecodeResult & result, + std::string & error) const; + + // Strict low-level API: every input ID must be an action token. Use + // decode_generated_tokens for complete Qwen sequences containing text. + bool decode_vlm_action_tokens(const std::vector> & batch_vlm_ids, FastDecodeResult & result, + std::string & error) const; + + // Production entry point for complete Qwen generated_ids. Ordinary text and + // control tokens are filtered through the explicit inverse action-token map. + bool decode_generated_tokens(const std::vector> & batch_generated_ids, + FastDecodeResult & result, std::string & error) const; + + private: + FastCodec(FastCodecConfig config, std::vector> token_bytes, + std::vector fast_to_vlm_id); + + bool byte_level_decode(const std::vector & fast_ids, std::vector & codepoints, + std::string & error) const; + + FastCodecConfig config_; + std::vector> token_bytes_; + std::vector fast_to_vlm_id_; + std::vector> vlm_to_fast_id_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_policy.cpp b/src/models/starvla/fast_policy.cpp new file mode 100644 index 0000000..e416fbb --- /dev/null +++ b/src/models/starvla/fast_policy.cpp @@ -0,0 +1,272 @@ +#include "models/starvla/fast_policy.h" + +#include "ggml.h" +#include "gguf.h" +#include "models/starvla/policy_gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kArchitecture = "starvla-policy"; +constexpr const char * kActionMapTensor = "starvla.policy.fast.action_token_map"; +constexpr const char * kOffsetsTensor = "starvla.policy.fast.codec.token_offsets"; +constexpr const char * kTokenBytesTensor = "starvla.policy.fast.codec.token_bytes"; + +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; + +struct FastRuntimeMetadata { + FastCodecConfig codec; + int token_bytes_count = 0; +}; + +FastRuntimeMetadata parse_metadata(gguf_context * gguf, FastPolicyConfig & config) { + if (require_string(gguf, "general.architecture") != kArchitecture || + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "fast") { + throw std::runtime_error("GGUF is not a supported StarVLA FAST policy"); + } + config.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config.text_filename = require_string(gguf, "starvla.component.text.filename"); + config.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config.backbone_arch != "qwen2_5_vl" || config.bundle_uuid.empty() || config.text_filename.empty() || + config.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA FAST bundle metadata is incomplete"); + } + + config.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config.action_dim = require_i32(gguf, "starvla.action.dimension"); + config.horizon = require_i32(gguf, "starvla.action.horizon"); + config.image_count = require_i32(gguf, "starvla.image.count"); + config.image_names = require_string_array(gguf, "starvla.image.names"); + config.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + const int max_length = require_i32(gguf, "starvla.fast.generation.max_length"); + config.generation_eos_token_ids = require_i32_array(gguf, "starvla.fast.generation.eos_token_ids"); + config.generation_top_k = require_i32(gguf, "starvla.fast.generation.top_k"); + config.generation_repetition_penalty = require_f32(gguf, "starvla.fast.generation.repetition_penalty"); + + FastRuntimeMetadata runtime; + runtime.codec.scale = require_f32(gguf, "starvla.fast.codec.scale"); + runtime.codec.min_token = require_i32(gguf, "starvla.fast.codec.min_token"); + runtime.codec.vocab_size = static_cast(require_i32(gguf, "starvla.fast.codec.vocab_size")); + runtime.codec.time_horizon = static_cast(require_i32(gguf, "starvla.fast.codec.time_horizon")); + runtime.codec.action_dim = static_cast(require_i32(gguf, "starvla.fast.codec.action_dimension")); + const int action_token_count = require_i32(gguf, "starvla.fast.action_token.count"); + const int offsets_count = require_i32(gguf, "starvla.fast.codec.token_offsets_count"); + runtime.token_bytes_count = require_i32(gguf, "starvla.fast.codec.token_bytes_count"); + + const bool valid = + config.qwen_hidden_dim > 0 && config.qwen_input_embedding_dim > 0 && config.qwen_vocab_size > 0 && + config.qwen_layer_count > 0 && !config.cot_template.empty() && config.action_dim > 0 && config.horizon > 0 && + config.image_count > 0 && config.image_names.size() == static_cast(config.image_count) && + config.image_processor_min_pixels > 0 && + config.image_processor_max_pixels >= config.image_processor_min_pixels && config.image_patch_size > 0 && + config.image_spatial_merge_size > 0 && config.image_min_token_count > 0 && + config.image_max_token_count >= config.image_min_token_count && max_length > 0 && + !config.generation_eos_token_ids.empty() && config.generation_top_k > 0 && + std::isfinite(config.generation_repetition_penalty) && config.generation_repetition_penalty > 0.0f && + runtime.codec.vocab_size > 0 && action_token_count == static_cast(runtime.codec.vocab_size) && + offsets_count == action_token_count + 1 && runtime.token_bytes_count > 0 && + runtime.codec.time_horizon == static_cast(config.horizon) && + runtime.codec.action_dim == static_cast(config.action_dim); + if (!valid) { + throw std::runtime_error("StarVLA FAST metadata has incompatible dimensions"); + } + config.generation_max_length = static_cast(max_length); + + config.normalization = detail::require_normalization(gguf, config.action_dim); + return runtime; +} + +struct RawTensor { + ggml_tensor * metadata = nullptr; + int index = -1; + std::vector bytes; +}; + +RawTensor read_tensor(const std::string & path, gguf_context * gguf, ggml_context * metadata_context, const char * name, + ggml_type expected_type, int64_t expected_elements) { + RawTensor result; + result.metadata = ggml_get_tensor(metadata_context, name); + result.index = gguf_find_tensor(gguf, name); + if (result.metadata == nullptr || result.index < 0 || result.metadata->type != expected_type || + ggml_n_dims(result.metadata) != 1 || result.metadata->ne[0] != expected_elements || + ggml_nelements(result.metadata) != expected_elements) { + throw std::runtime_error(std::string("FAST runtime tensor shape/type mismatch: ") + name); + } + result.bytes.resize(ggml_nbytes(result.metadata)); + std::ifstream stream(path, std::ios::binary); + if (!stream) { + throw std::runtime_error("failed to open FAST policy GGUF tensor data"); + } + const size_t offset = gguf_get_data_offset(gguf) + gguf_get_tensor_offset(gguf, result.index); + stream.seekg(static_cast(offset), std::ios::beg); + if (!stream || offset > static_cast(std::numeric_limits::max())) { + throw std::runtime_error(std::string("failed to seek FAST runtime tensor: ") + name); + } + stream.read(reinterpret_cast(result.bytes.data()), static_cast(result.bytes.size())); + if (!stream) { + throw std::runtime_error(std::string("failed to read FAST runtime tensor: ") + name); + } + return result; +} + +} // namespace + +struct FastPolicy::Impl { + FastPolicyConfig config; + std::unique_ptr codec; +}; + +FastPolicy::FastPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +FastPolicy::~FastPolicy() = default; + +std::unique_ptr FastPolicy::load(const std::string & path, int verbosity, std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA FAST policy path is required"; + return nullptr; + } + + ggml_context * metadata_context = nullptr; + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = &metadata_context; + gguf_context * gguf = gguf_init_from_file(path.c_str(), params); + if (gguf == nullptr || metadata_context == nullptr) { + if (metadata_context != nullptr) { + ggml_free(metadata_context); + } + if (gguf != nullptr) { + gguf_free(gguf); + } + error = "failed to read StarVLA FAST policy GGUF"; + return nullptr; + } + auto cleanup = [&]() { + ggml_free(metadata_context); + metadata_context = nullptr; + gguf_free(gguf); + gguf = nullptr; + }; + + std::unique_ptr impl(new Impl()); + try { + const FastRuntimeMetadata runtime = parse_metadata(gguf, impl->config); + + RawTensor action_map = read_tensor(path, gguf, metadata_context, kActionMapTensor, GGML_TYPE_I32, + static_cast(runtime.codec.vocab_size)); + RawTensor offsets = read_tensor(path, gguf, metadata_context, kOffsetsTensor, GGML_TYPE_I32, + static_cast(runtime.codec.vocab_size + 1)); + RawTensor token_bytes = + read_tensor(path, gguf, metadata_context, kTokenBytesTensor, GGML_TYPE_I8, runtime.token_bytes_count); + + const uint32_t endian_probe = 1; + if (*reinterpret_cast(&endian_probe) != 1) { + throw std::runtime_error("FAST runtime currently requires a little-endian host"); + } + std::vector action_ids(runtime.codec.vocab_size); + std::vector token_offsets(runtime.codec.vocab_size + 1); + std::memcpy(action_ids.data(), action_map.bytes.data(), action_map.bytes.size()); + std::memcpy(token_offsets.data(), offsets.bytes.data(), offsets.bytes.size()); + impl->codec = FastCodec::create_compiled(runtime.codec, std::move(token_offsets), std::move(token_bytes.bytes), + std::move(action_ids), error); + if (impl->codec == nullptr) { + throw std::runtime_error("failed to construct embedded FAST codec: " + error); + } + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: bundle=%s runtime_tensors=3 codec_vocab=%zu " + "generation_max_length=%zu profiles=%zu\n", + __func__, impl->config.bundle_uuid.c_str(), runtime.codec.vocab_size, + impl->config.generation_max_length, impl->config.normalization.profiles.size()); + } + cleanup(); + } catch (const std::exception & exception) { + cleanup(); + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new FastPolicy(std::move(impl))); +} + +bool FastPolicy::decode_generated(const std::vector & full_sequence, std::vector & normalized_actions, + std::string & error) const { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->codec == nullptr) { + error = "StarVLA FAST policy is not initialized"; + return false; + } + FastDecodeResult decoded; + if (!impl_->codec->decode_generated_tokens({full_sequence}, decoded, error)) { + return false; + } + if (decoded.batch_size != 1 || decoded.time_horizon != static_cast(impl_->config.horizon) || + decoded.action_dim != static_cast(impl_->config.action_dim) || + decoded.actions.size() != static_cast(impl_->config.horizon * impl_->config.action_dim)) { + error = "embedded FAST codec returned an incompatible action tensor"; + return false; + } + normalized_actions.reserve(decoded.actions.size()); + for (double value : decoded.actions) { + const float converted = static_cast(value); + if (!std::isfinite(converted)) { + normalized_actions.clear(); + error = "embedded FAST codec returned a non-finite action"; + return false; + } + normalized_actions.push_back(converted); + } + return true; +} + +bool FastPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA FAST policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key, normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const FastPolicyConfig & FastPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA FAST policy is not initialized"); + } + return impl_->config; +} + +const char * FastPolicy::backend_name() const { + return "cpu"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/fast_policy.h b/src/models/starvla/fast_policy.h new file mode 100644 index 0000000..a92bbfb --- /dev/null +++ b/src/models/starvla/fast_policy.h @@ -0,0 +1,72 @@ +#pragma once + +#include "models/starvla/fast_codec.h" +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct FastPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_vocab_size = 0; + int qwen_layer_count = 0; + + std::string cot_template; + int action_dim = 0; + int horizon = 0; + + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + size_t generation_max_length = 0; + std::vector generation_eos_token_ids; + int generation_top_k = 0; + float generation_repetition_penalty = 0.0f; + + NormalizationConfig normalization; +}; + +class FastPolicy { + public: + ~FastPolicy(); + + FastPolicy(const FastPolicy &) = delete; + FastPolicy & operator=(const FastPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int verbosity, std::string & error); + + bool decode_generated(const std::vector & full_sequence, std::vector & normalized_actions, + std::string & error) const; + + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const FastPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit FastPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_policy.cpp b/src/models/starvla/groot_policy.cpp new file mode 100644 index 0000000..a91a150 --- /dev/null +++ b/src/models/starvla/groot_policy.cpp @@ -0,0 +1,693 @@ +#include "models/starvla/groot_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 16384; +constexpr int kKQMaskPad = 32; + +struct GR00TBlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct GR00TWeights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + ggml_tensor * output_modulation_weight = nullptr; + ggml_tensor * output_modulation_bias = nullptr; + ggml_tensor * output_projection_weight = nullptr; + ggml_tensor * output_projection_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; + +class GR00TGGUFLoader final : public gguf_loader { + public: + GR00TGGUFLoader(GR00TPolicyConfig & config, GR00TWeights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy") { + throw std::runtime_error("StarVLA GR00T policy has incompatible general.architecture"); + } + if (require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "groot") { + throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen GR00T schema"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl" && config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error("StarVLA GR00T policy has an unsupported Qwen backbone"); + } + + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + if (config_.bundle_uuid.empty()) { + throw std::runtime_error("StarVLA GR00T bundle UUID is missing"); + } + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA GR00T component filenames must be non-empty"); + } + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + const bool qwen25 = config_.backbone_arch == "qwen2_5_vl"; + if (config_.cot_template.empty()) { + throw std::runtime_error("StarVLA GR00T prompt template is missing"); + } + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + config_.dit_width = require_i32(gguf, "starvla.groot.dit_width"); + config_.block_count = require_i32(gguf, "starvla.groot.block_count"); + config_.attention_head_count = require_i32(gguf, "starvla.groot.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.groot.attention_head_dim"); + config_.cross_attention_dim = require_i32(gguf, "starvla.groot.cross_attention_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.groot.feed_forward_dim"); + config_.output_dim = require_i32(gguf, "starvla.groot.output_dimension"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.groot.mlp_hidden_dimension"); + config_.future_token_count = require_i32(gguf, "starvla.groot.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.groot.action_position_count"); + config_.no_state_sequence_length = require_i32(gguf, "starvla.groot.no_state_sequence_length"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.groot.timestep_projection_dim"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.groot.ada_norm_epsilon"); + config_.output_norm_epsilon = require_f32(gguf, "starvla.groot.output_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.groot.euler_dt"); + config_.timestep_ids = require_i32_array(gguf, "starvla.groot.timestep_ids"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + + const int64_t expected_input_embedding_dim = + qwen25 ? static_cast(config_.qwen_hidden_dim) : 4LL * config_.qwen_hidden_dim; + const bool dimensions_valid = + config_.qwen_hidden_dim > 0 && config_.qwen_vocab_size > 0 && config_.dit_width > 0 && + config_.qwen_input_embedding_dim == expected_input_embedding_dim && config_.dit_width % 2 == 0 && + config_.block_count > 0 && config_.block_count % 2 == 0 && config_.attention_head_count > 0 && + config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.cross_attention_dim == config_.qwen_hidden_dim && config_.feed_forward_dim > 0 && + config_.output_dim > 0 && config_.mlp_hidden_dim > 0 && config_.action_dim > 0 && config_.horizon > 0 && + config_.future_token_count > 0 && config_.action_position_count >= config_.horizon && + config_.no_state_sequence_length == config_.future_token_count + config_.horizon && + config_.timestep_projection_dim >= 4 && config_.timestep_projection_dim % 2 == 0 && + std::isfinite(config_.ada_norm_epsilon) && config_.ada_norm_epsilon > 0.0f && + std::isfinite(config_.output_norm_epsilon) && config_.output_norm_epsilon > 0.0f && + std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && config_.timestep_ids.size() == 4 && + config_.image_count > 0 && config_.image_names.size() == static_cast(config_.image_count) && + config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count; + if (!dimensions_valid) { + throw std::runtime_error("StarVLA GR00T policy metadata has incompatible dimensions"); + } + config_.normalization = detail::require_normalization(gguf, config_.action_dim); + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, "starvla.policy.groot.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.groot.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.groot.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.groot.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.groot.block." + std::to_string(block) + "."; + GR00TBlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + bind(weights_.output_modulation_weight, "starvla.policy.groot.output.modulation.weight"); + bind(weights_.output_modulation_bias, "starvla.policy.groot.output.modulation.bias"); + bind(weights_.output_projection_weight, "starvla.policy.groot.output.projection.weight"); + bind(weights_.output_projection_bias, "starvla.policy.groot.output.projection.bias"); + bind(weights_.action_input_weight, "starvla.policy.groot.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.groot.action.input.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.groot.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.groot.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.groot.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.groot.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.groot.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.groot.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.groot.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.groot.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.groot.future_tokens.weight"); + bind(weights_.action_position, "starvla.policy.groot.action_position.weight"); + + const int width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.output_modulation_weight, {width, 2 * width}) || + !has_shape(weights_.output_modulation_bias, {2 * width}) || + !has_shape(weights_.output_projection_weight, {width, config_.output_dim}) || + !has_shape(weights_.output_projection_bias, {config_.output_dim}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, {config_.output_dim, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA GR00T non-block tensor has an incompatible ggml shape"); + } + for (int block = 0; block < config_.block_count; ++block) { + const GR00TBlockWeights & current = weights_.blocks[static_cast(block)]; + const int kv_input_dim = block % 2 == 0 ? config_.cross_attention_dim : width; + if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || !has_shape(current.key_weight, {kv_input_dim, width}) || + !has_shape(current.key_bias, {width}) || !has_shape(current.value_weight, {kv_input_dim, width}) || + !has_shape(current.value_bias, {width}) || + !has_shape(current.attention_output_weight, {width, width}) || + !has_shape(current.attention_output_bias, {width}) || + !has_shape(current.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(current.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(current.feed_forward_output_weight, {config_.feed_forward_dim, width}) || + !has_shape(current.feed_forward_output_bias, {width})) { + throw std::runtime_error("StarVLA GR00T transformer block tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + GR00TPolicyConfig & config_; + GR00TWeights & weights_; +}; + +std::vector timestep_projection_table(const GR00TPolicyConfig & config) { + const int dim = config.timestep_projection_dim; + const int half = dim / 2; + const float denominator = static_cast(half - 1); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); + } + } + return table; +} + +std::vector action_time_table(const GR00TPolicyConfig & config) { + const int dim = config.dit_width; + const int half = dim / 2; + const float denominator = static_cast(half); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); + } + } + return table; +} + +} // namespace + +struct GR00TPolicy::Impl { + GR00TPolicyConfig config; + GR00TWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + std::vector timestep_table; + std::vector action_table; + + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + } + + void build_graph(size_t token_count) { + clear_graph(); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA GR00T conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA GR00T graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + const int sequence_length = config.no_state_sequence_length; + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + + hidden_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, static_cast(token_count)); + cross_mask_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA GR00T graph inputs"); + } + ggml_set_name(hidden_input, "starvla_groot_qwen_hidden_states"); + ggml_set_name(cross_mask_input, "starvla_groot_qwen_attention_mask"); + ggml_set_name(noise_input, "starvla_groot_initial_noise"); + ggml_set_name(timestep_projection_input, "starvla_groot_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_groot_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_groot_scalar_one"); + ggml_set_input(hidden_input); + ggml_set_input(cross_mask_input); + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const GR00TBlockWeights & block) { + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + ggml_tensor * one_plus_scale = ggml_add(graph_context, scale, scalar_one_input); + return ggml_add(graph_context, ggml_mul(graph_context, normalized, one_plus_scale), shift); + }; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, ggml_tensor * mask, + const GR00TBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = + ggml_soft_max_ext(graph_context, scores, mask, 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); + }; + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; + + for (int step = 0; step < 4; ++step) { + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); + ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const GR00TBlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = block_index % 2 == 0 + ? attention(normalized, hidden_input, cross_mask_input, block) + : attention(normalized, normalized, nullptr, block); + hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + ggml_tensor * output_modulation = linear(ggml_silu(graph_context, temb), weights.output_modulation_weight, + weights.output_modulation_bias); + // DiT output uses shift then scale, unlike AdaLayerNorm's scale then shift. + ggml_tensor * shift = ggml_view_1d(graph_context, output_modulation, width, 0); + ggml_tensor * scale = + ggml_view_1d(graph_context, output_modulation, width, static_cast(width) * sizeof(float)); + hidden = ggml_norm(graph_context, hidden, config.output_norm_epsilon); + hidden = ggml_mul(graph_context, hidden, ggml_add(graph_context, scale, scalar_one_input)); + hidden = ggml_add(graph_context, hidden, shift); + hidden = linear(hidden, weights.output_projection_weight, weights.output_projection_bias); + hidden = + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); + ggml_tensor * velocity = + ggml_view_2d(graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_groot_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA GR00T graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA GR00T graph"); + } + + conditioning_token_count = token_count; + } +}; + +GR00TPolicy::GR00TPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +GR00TPolicy::~GR00TPolicy() = default; + +std::unique_ptr GR00TPolicy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA GR00T policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { + error = "failed to initialize StarVLA GR00T backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + GR00TGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA GR00T policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf(stderr, "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new GR00TPolicy(std::move(impl))); +} + +bool GR00TPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA GR00T policy is not initialized"; + return false; + } + if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || + mask_element_count == 0 || mask_element_count > static_cast(std::numeric_limits::max()) || + mask_element_count > std::numeric_limits::max() / static_cast(impl_->config.qwen_hidden_dim) || + hidden_element_count != mask_element_count * static_cast(impl_->config.qwen_hidden_dim)) { + error = "StarVLA GR00T Qwen conditioning tensor or attention mask has an incompatible shape"; + return false; + } + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (noise_element_count != expected_noise) { + error = "StarVLA GR00T initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = "StarVLA GR00T conditioning and initial noise must be finite"; + return false; + } + bool has_valid_token = false; + for (size_t token = 0; token < mask_element_count; ++token) { + if (qwen_attention_mask[token] > 1) { + error = "StarVLA GR00T attention mask values must be zero or one"; + return false; + } + has_valid_token = has_valid_token || qwen_attention_mask[token] != 0; + } + if (!has_valid_token) { + error = "StarVLA GR00T attention mask must contain at least one valid token"; + return false; + } + + try { + if (impl_->graph == nullptr || impl_->conditioning_token_count != mask_element_count) { + impl_->build_graph(mask_element_count); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + const int query_count = impl_->config.no_state_sequence_length; + const int padded_queries = GGML_PAD(query_count, kKQMaskPad); + std::vector additive_mask(mask_element_count * static_cast(padded_queries), + -std::numeric_limits::infinity()); + for (int query = 0; query < query_count; ++query) { + float * row = additive_mask.data() + static_cast(query) * mask_element_count; + for (size_t token = 0; token < mask_element_count; ++token) { + row[token] = qwen_attention_mask[token] != 0 ? 0.0f : -std::numeric_limits::infinity(); + } + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA GR00T graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA GR00T graph produced non-finite actions"; + return false; + } + return true; +} + +bool GR00TPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA GR00T policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); +} + +const GR00TPolicyConfig & GR00TPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA GR00T policy is not initialized"); + } + return impl_->config; +} + +const char * GR00TPolicy::backend_name() const { + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_policy.h b/src/models/starvla/groot_policy.h new file mode 100644 index 0000000..61cc2cc --- /dev/null +++ b/src/models/starvla/groot_policy.h @@ -0,0 +1,81 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct GR00TPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_vocab_size = 0; + std::string cot_template; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int output_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; + int no_state_sequence_length = 0; + int timestep_projection_dim = 0; + float ada_norm_epsilon = 0.0f; + float output_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class GR00TPolicy { + public: + ~GR00TPolicy(); + + GR00TPolicy(const GR00TPolicy &) = delete; + GR00TPolicy & operator=(const GR00TPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, + std::string & error); + + // qwen_hidden_states is token-major [token_count, qwen_hidden_dim]. The mask + // follows torch SDPA semantics: non-zero entries participate in attention. + // initial_noise is token-major [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const uint8_t * qwen_attention_mask, + size_t mask_element_count, const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const GR00TPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit GR00TPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_prompt.cpp b/src/models/starvla/groot_prompt.cpp new file mode 100644 index 0000000..7b490aa --- /dev/null +++ b/src/models/starvla/groot_prompt.cpp @@ -0,0 +1,68 @@ +#include "models/starvla/groot_prompt.h" + +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kInstructionPlaceholder = "{instruction}"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; + +bool contains_nul(const std::string & value) { + return value.find('\0') != std::string::npos; +} + +void replace_all(std::string & value, const std::string & needle, const std::string & replacement) { + size_t offset = 0; + while ((offset = value.find(needle, offset)) != std::string::npos) { + value.replace(offset, needle.size(), replacement); + offset += replacement.size(); + } +} + +bool build_instruction(const char * framework, const std::string & cot_template, const std::string & task, + std::string & instruction, std::string & error) { + instruction.clear(); + error.clear(); + + if (task.empty()) { + error = std::string("StarVLA ") + framework + " task must not be empty"; + return false; + } + if (contains_nul(task) || contains_nul(cot_template)) { + error = std::string("StarVLA ") + framework + " prompt contains an embedded NUL byte"; + return false; + } + if (task.find(kMtmdMediaMarker) != std::string::npos || cot_template.find(kMtmdMediaMarker) != std::string::npos) { + error = std::string("StarVLA ") + framework + " prompt contains the reserved mtmd media marker"; + return false; + } + if (cot_template.find(kInstructionPlaceholder) == std::string::npos) { + error = std::string("StarVLA ") + framework + " CoT template is missing {instruction}"; + return false; + } + + std::string wrapped = cot_template; + replace_all(wrapped, kInstructionPlaceholder, task); + instruction = std::move(wrapped); + return true; +} + +} // namespace + +bool build_groot_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { + return build_instruction("GR00T", cot_template, task, instruction, error); +} + +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { + return build_instruction("PI_v3", cot_template, task, instruction, error); +} + +bool build_fast_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error) { + return build_instruction("FAST", cot_template, task, instruction, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/groot_prompt.h b/src/models/starvla/groot_prompt.h new file mode 100644 index 0000000..d14eee2 --- /dev/null +++ b/src/models/starvla/groot_prompt.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace robotcpp::starvla { + +bool build_groot_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); + +bool build_pi_v3_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); + +bool build_fast_instruction(const std::string & cot_template, const std::string & task, std::string & instruction, + std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/normalization.cpp b/src/models/starvla/normalization.cpp new file mode 100644 index 0000000..8fd03a4 --- /dev/null +++ b/src/models/starvla/normalization.cpp @@ -0,0 +1,170 @@ +#include "models/starvla/normalization.h" + +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +std::string profile_keys(const NormalizationConfig & config) { + std::ostringstream out; + for (size_t i = 0; i < config.profiles.size(); ++i) { + if (i != 0) { + out << ", "; + } + out << config.profiles[i].key; + } + return out.str(); +} + +} // namespace + +bool validate_normalization_config(const NormalizationConfig & config, int action_dim, std::string & error) { + error.clear(); + if (action_dim <= 0) { + error = "StarVLA normalization requires a positive action dimension"; + return false; + } + if (!std::isfinite(config.binary_threshold) || config.binary_threshold != 0.5f) { + error = "StarVLA normalization binary threshold must use the canonical value 0.5"; + return false; + } + if (config.binary_comparison != "gt" && config.binary_comparison != "ge") { + error = "StarVLA normalization binary comparison must be 'gt' or 'ge'"; + return false; + } + if (config.profiles.empty()) { + error = "StarVLA policy has no normalization profiles"; + return false; + } + if (config.default_profile_key.empty()) { + error = "StarVLA policy has no default normalization profile"; + return false; + } + + std::vector dimension_kind(static_cast(action_dim), 0); + for (int32_t dim : config.continuous_dimensions) { + if (dim < 0 || dim >= action_dim || dimension_kind[static_cast(dim)] != 0) { + error = "StarVLA continuous action dimensions are invalid or duplicated"; + return false; + } + dimension_kind[static_cast(dim)] = 1; + } + for (int32_t dim : config.binary_dimensions) { + if (dim < 0 || dim >= action_dim || dimension_kind[static_cast(dim)] != 0) { + error = "StarVLA binary action dimensions are invalid or duplicated"; + return false; + } + dimension_kind[static_cast(dim)] = 2; + } + if (std::find(dimension_kind.begin(), dimension_kind.end(), uint8_t{0}) != dimension_kind.end()) { + error = "StarVLA continuous and binary action dimensions must cover every action column"; + return false; + } + + std::vector seen_keys; + seen_keys.reserve(config.profiles.size()); + for (const NormalizationProfile & profile : config.profiles) { + if (profile.key.empty() || std::find(seen_keys.begin(), seen_keys.end(), profile.key) != seen_keys.end()) { + error = "StarVLA normalization profile keys must be non-empty and unique"; + return false; + } + seen_keys.push_back(profile.key); + if (profile.action_q01.size() != static_cast(action_dim) || + profile.action_q99.size() != static_cast(action_dim) || + profile.action_mask.size() != static_cast(action_dim)) { + error = "StarVLA normalization profile shape does not match action dimension: " + profile.key; + return false; + } + for (int dim = 0; dim < action_dim; ++dim) { + const size_t index = static_cast(dim); + if (!std::isfinite(profile.action_q01[index]) || !std::isfinite(profile.action_q99[index])) { + error = "StarVLA normalization quantiles must be finite: " + profile.key; + return false; + } + if (dimension_kind[index] == 1 && profile.action_q99[index] < profile.action_q01[index]) { + error = "StarVLA normalization q99 must not be below q01: " + profile.key; + return false; + } + if (dimension_kind[index] == 1 && profile.action_mask[index] == 0) { + error = "StarVLA continuous action dimension is disabled by the normalization mask: " + profile.key; + return false; + } + if (dimension_kind[index] == 2 && profile.action_mask[index] != 0) { + error = "StarVLA binary action dimension must not use q01/q99 scaling: " + profile.key; + return false; + } + } + } + if (std::find(seen_keys.begin(), seen_keys.end(), config.default_profile_key) == seen_keys.end()) { + error = "StarVLA default normalization profile is not present: " + config.default_profile_key; + return false; + } + return true; +} + +const NormalizationProfile * resolve_normalization_profile(const NormalizationConfig & config, + const std::string & profile_key, std::string & error) { + error.clear(); + if (profile_key.empty()) { + return resolve_normalization_profile(config, config.default_profile_key, error); + } + for (const NormalizationProfile & profile : config.profiles) { + if (profile.key == profile_key) { + return &profile; + } + } + error = "unknown StarVLA normalization profile '" + profile_key + "'; expected one of: " + profile_keys(config); + return nullptr; +} + +bool denormalize_actions(const NormalizationConfig & config, const std::string & profile_key, + const std::vector & normalized, int horizon, int action_dim, + std::vector & actions, std::string & error) { + actions.clear(); + error.clear(); + if (!validate_normalization_config(config, action_dim, error)) { + return false; + } + if (horizon <= 0 || normalized.size() != static_cast(horizon) * static_cast(action_dim)) { + error = "StarVLA normalized action tensor has an incompatible shape"; + return false; + } + const NormalizationProfile * profile = resolve_normalization_profile(config, profile_key, error); + if (profile == nullptr) { + return false; + } + + std::vector is_binary(static_cast(action_dim), 0); + for (int32_t dim : config.binary_dimensions) { + is_binary[static_cast(dim)] = 1; + } + + actions.resize(normalized.size()); + for (int step = 0; step < horizon; ++step) { + for (int dim = 0; dim < action_dim; ++dim) { + const size_t index = static_cast(step) * static_cast(action_dim) + static_cast(dim); + const float input_value = normalized[index]; + if (!std::isfinite(input_value)) { + actions.clear(); + error = "StarVLA normalized actions must be finite"; + return false; + } + const float value = config.clip_actions ? std::clamp(input_value, -1.0f, 1.0f) : input_value; + if (is_binary[static_cast(dim)] != 0) { + const bool active = config.binary_comparison == "ge" ? value >= config.binary_threshold + : value > config.binary_threshold; + actions[index] = active ? 1.0f : 0.0f; + } else { + const float low = profile->action_q01[static_cast(dim)]; + const float high = profile->action_q99[static_cast(dim)]; + actions[index] = (value + 1.0f) * 0.5f * (high - low) + low; + } + } + } + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/normalization.h b/src/models/starvla/normalization.h new file mode 100644 index 0000000..696830b --- /dev/null +++ b/src/models/starvla/normalization.h @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include + +namespace robotcpp::starvla { + +struct NormalizationProfile { + std::string key; + std::vector action_q01; + std::vector action_q99; + std::vector action_mask; +}; + +struct NormalizationConfig { + std::string default_profile_key; + bool clip_actions = false; + float binary_threshold = 0.5f; + std::string binary_comparison; + std::vector continuous_dimensions; + std::vector binary_dimensions; + std::vector profiles; +}; + +bool validate_normalization_config(const NormalizationConfig & config, int action_dim, std::string & error); + +const NormalizationProfile * resolve_normalization_profile(const NormalizationConfig & config, + const std::string & profile_key, std::string & error); + +bool denormalize_actions(const NormalizationConfig & config, const std::string & profile_key, + const std::vector & normalized, int horizon, int action_dim, + std::vector & actions, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_image_preprocess.cpp b/src/models/starvla/oft_image_preprocess.cpp new file mode 100644 index 0000000..ecace78 --- /dev/null +++ b/src/models/starvla/oft_image_preprocess.cpp @@ -0,0 +1,339 @@ +#include "models/starvla/oft_image_preprocess.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +struct RGBImage { + int width = 0; + int height = 0; + std::vector pixels; +}; + +struct FilterTable { + int kernel_size = 0; + int precision = 0; + std::vector first; + std::vector count; + std::vector weights; +}; + +double keys_cubic(double value) { + constexpr double a = -0.5; + value = std::fabs(value); + if (value < 1.0) { + return ((a + 2.0) * value - (a + 3.0)) * value * value + 1.0; + } + if (value < 2.0) { + return (((value - 5.0) * value + 8.0) * value - 4.0) * a; + } + return 0.0; +} + +bool validate_resize(const uint8_t * source, int source_width, int source_height, int source_stride, int target_width, + int target_height, std::string & error) { + error.clear(); + if (source == nullptr || source_width <= 0 || source_height <= 0 || target_width <= 0 || target_height <= 0) { + error = "StarVLA image resize received an invalid image or dimension"; + return false; + } + const int packed_stride = source_width * 3; + if (source_stride != 0 && source_stride < packed_stride) { + error = "StarVLA image stride is smaller than a packed RGB row"; + return false; + } + const uint64_t output_bytes = static_cast(target_width) * static_cast(target_height) * 3; + if (output_bytes > static_cast(std::numeric_limits::max())) { + error = "StarVLA resized image is too large"; + return false; + } + return true; +} + +RGBImage pack_source(const uint8_t * source, int width, int height, int stride) { + RGBImage image; + image.width = width; + image.height = height; + const size_t row_bytes = static_cast(width) * 3; + const size_t actual_stride = stride > 0 ? static_cast(stride) : row_bytes; + image.pixels.resize(row_bytes * static_cast(height)); + for (int row = 0; row < height; ++row) { + std::copy(source + static_cast(row) * actual_stride, + source + static_cast(row) * actual_stride + row_bytes, + image.pixels.begin() + static_cast(static_cast(row) * row_bytes)); + } + return image; +} + +FilterTable make_filter_table(int input_size, int output_size, bool pillow_precision) { + const double scale = static_cast(input_size) / static_cast(output_size); + const double filter_scale = std::max(scale, 1.0); + const double support = 2.0 * filter_scale; + + FilterTable table; + table.kernel_size = static_cast(std::ceil(support)) * 2 + 1; + table.first.resize(static_cast(output_size)); + table.count.resize(static_cast(output_size)); + std::vector floating_weights(static_cast(output_size) * static_cast(table.kernel_size), + 0.0); + double maximum_weight = 0.0; + + for (int output = 0; output < output_size; ++output) { + const double center = (static_cast(output) + 0.5) * scale; + const int first = std::max(static_cast(center - support + 0.5), 0); + const int end = std::min(static_cast(center + support + 0.5), input_size); + const int count = std::max(0, std::min(end - first, table.kernel_size)); + table.first[static_cast(output)] = first; + table.count[static_cast(output)] = count; + + double sum = 0.0; + for (int index = 0; index < count; ++index) { + const double distance = (static_cast(index + first) - center + 0.5) / filter_scale; + const double weight = keys_cubic(distance); + floating_weights[static_cast(output) * table.kernel_size + index] = weight; + sum += weight; + } + if (sum != 0.0) { + for (int index = 0; index < count; ++index) { + double & weight = floating_weights[static_cast(output) * table.kernel_size + index]; + weight /= sum; + maximum_weight = std::max(maximum_weight, weight); + } + } + } + + if (pillow_precision) { + table.precision = 22; + } else { + for (table.precision = 0; table.precision < 22; ++table.precision) { + const int next = + static_cast(0.5 + maximum_weight * static_cast(uint32_t{1} << (table.precision + 1))); + if (next >= (1 << 15)) { + break; + } + } + } + + const double multiplier = static_cast(uint32_t{1} << table.precision); + table.weights.resize(floating_weights.size()); + for (size_t index = 0; index < floating_weights.size(); ++index) { + const double scaled = floating_weights[index] * multiplier; + table.weights[index] = static_cast(scaled < 0.0 ? scaled - 0.5 : scaled + 0.5); + } + return table; +} + +uint8_t fixed_point_pixel(int64_t accumulator, int precision) { + const int64_t value = accumulator >> precision; + return static_cast(std::max(0, std::min(255, value))); +} + +RGBImage resize_horizontal(const RGBImage & source, int target_width, const FilterTable & table) { + RGBImage target; + target.width = target_width; + target.height = source.height; + target.pixels.resize(static_cast(target.width) * target.height * 3); + const int64_t rounding = int64_t{1} << (table.precision - 1); + for (int row = 0; row < source.height; ++row) { + for (int column = 0; column < target.width; ++column) { + const int first = table.first[static_cast(column)]; + const int count = table.count[static_cast(column)]; + for (int channel = 0; channel < 3; ++channel) { + int64_t accumulator = rounding; + for (int index = 0; index < count; ++index) { + const size_t source_index = (static_cast(row) * source.width + first + index) * 3 + channel; + const int32_t weight = table.weights[static_cast(column) * table.kernel_size + index]; + accumulator += static_cast(source.pixels[source_index]) * weight; + } + const size_t target_index = (static_cast(row) * target.width + column) * 3 + channel; + target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); + } + } + } + return target; +} + +RGBImage resize_vertical(const RGBImage & source, int target_height, const FilterTable & table) { + RGBImage target; + target.width = source.width; + target.height = target_height; + target.pixels.resize(static_cast(target.width) * target.height * 3); + const int64_t rounding = int64_t{1} << (table.precision - 1); + for (int row = 0; row < target.height; ++row) { + const int first = table.first[static_cast(row)]; + const int count = table.count[static_cast(row)]; + for (int column = 0; column < target.width; ++column) { + for (int channel = 0; channel < 3; ++channel) { + int64_t accumulator = rounding; + for (int index = 0; index < count; ++index) { + const size_t source_index = + (static_cast(first + index) * source.width + column) * 3 + channel; + const int32_t weight = table.weights[static_cast(row) * table.kernel_size + index]; + accumulator += static_cast(source.pixels[source_index]) * weight; + } + const size_t target_index = (static_cast(row) * target.width + column) * 3 + channel; + target.pixels[target_index] = fixed_point_pixel(accumulator, table.precision); + } + } + } + return target; +} + +bool resize_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, int target_width, + int target_height, bool pillow_precision, std::vector & target, std::string & error) { + target.clear(); + if (!validate_resize(source, source_width, source_height, source_stride, target_width, target_height, error)) { + return false; + } + + RGBImage current = pack_source(source, source_width, source_height, source_stride); + if (source_width != target_width) { + current = + resize_horizontal(current, target_width, make_filter_table(source_width, target_width, pillow_precision)); + } + if (source_height != target_height) { + current = + resize_vertical(current, target_height, make_filter_table(source_height, target_height, pillow_precision)); + } + target = std::move(current.pixels); + return true; +} + +} // namespace + +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, + std::string & error) { + return resize_rgb(source, source_width, source_height, source_stride, target_width, target_height, true, target, + error); +} + +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, + std::string & error) { + return resize_rgb(source, source_width, source_height, source_stride, target_width, target_height, false, target, + error); +} + +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, int min_pixels, int max_pixels, + int & target_width, int & target_height, std::string & error) { + target_width = 0; + target_height = 0; + error.clear(); + if (source_width <= 0 || source_height <= 0 || factor <= 0 || min_pixels <= 0 || max_pixels < min_pixels) { + error = "Qwen3-VL smart resize received an invalid dimension or pixel bound"; + return false; + } + const int minimum_side = std::min(source_width, source_height); + const int maximum_side = std::max(source_width, source_height); + if (static_cast(maximum_side) / minimum_side > 200.0) { + error = "Qwen3-VL smart resize requires an absolute aspect ratio of at most 200"; + return false; + } + + // Python round() uses ties-to-even. Integer quotient/remainder arithmetic + // makes the common first smart_resize step independent of the host FP mode. + const auto round_div_ties_to_even = [](int value, int divisor) -> int64_t { + const int64_t quotient = value / divisor; + const int64_t remainder = value % divisor; + const int64_t doubled = remainder * 2; + if (doubled < divisor || (doubled == divisor && quotient % 2 == 0)) { + return quotient; + } + return quotient + 1; + }; + + int64_t resized_height = round_div_ties_to_even(source_height, factor) * factor; + int64_t resized_width = round_div_ties_to_even(source_width, factor) * factor; + const int64_t source_pixels = static_cast(source_height) * source_width; + const int64_t rounded_pixels = resized_height * resized_width; + if (rounded_pixels > max_pixels) { + const double beta = std::sqrt(static_cast(source_pixels) / max_pixels); + resized_height = + std::max(factor, static_cast(std::floor(source_height / beta / factor)) * factor); + resized_width = + std::max(factor, static_cast(std::floor(source_width / beta / factor)) * factor); + } else if (rounded_pixels < min_pixels) { + const double beta = std::sqrt(static_cast(min_pixels) / source_pixels); + resized_height = static_cast(std::ceil(source_height * beta / factor)) * factor; + resized_width = static_cast(std::ceil(source_width * beta / factor)) * factor; + } + + if (resized_width <= 0 || resized_height <= 0 || resized_width > std::numeric_limits::max() || + resized_height > std::numeric_limits::max() || + resized_width > std::numeric_limits::max() / resized_height) { + error = "Qwen3-VL smart resize produced an unsupported output dimension"; + return false; + } + target_width = static_cast(resized_width); + target_height = static_cast(resized_height); + return true; +} + +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, int channels, + int source_stride, int patch_size, int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, int & target_height, + int & image_token_count, std::string & error) { + target.clear(); + target_width = 0; + target_height = 0; + image_token_count = 0; + if (channels != 3) { + error = "Qwen3-VL input image must be RGB"; + return false; + } + if (patch_size <= 0 || spatial_merge_size <= 0 || + patch_size > std::numeric_limits::max() / spatial_merge_size) { + error = "Qwen3-VL patch or spatial merge size is invalid"; + return false; + } + const int factor = patch_size * spatial_merge_size; + if (!qwen3vl_smart_resize_dimensions(source_width, source_height, factor, min_pixels, max_pixels, target_width, + target_height, error)) { + return false; + } + if (!resize_torchvision_bicubic_aa_rgb(source, source_width, source_height, source_stride, target_width, + target_height, target, error)) { + target_width = 0; + target_height = 0; + return false; + } + const int64_t grid_width = target_width / factor; + const int64_t grid_height = target_height / factor; + const int64_t tokens = grid_width * grid_height; + if (tokens <= 0 || tokens > std::numeric_limits::max()) { + target.clear(); + target_width = 0; + target_height = 0; + error = "Qwen3-VL smart resize produced an unsupported image token count"; + return false; + } + image_token_count = static_cast(tokens); + return true; +} + +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, int channels, int source_stride, + int training_width, int training_height, int processor_width, int processor_height, + std::vector & target, std::string & error) { + target.clear(); + if (channels != 3) { + error = "StarVLA OFT input image must be RGB"; + return false; + } + std::vector training_image; + if (!resize_pillow_bicubic_rgb(source, source_width, source_height, source_stride, training_width, training_height, + training_image, error)) { + return false; + } + return resize_torchvision_bicubic_aa_rgb(training_image.data(), training_width, training_height, training_width * 3, + processor_width, processor_height, target, error); +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_image_preprocess.h b/src/models/starvla/oft_image_preprocess.h new file mode 100644 index 0000000..6aaa2a7 --- /dev/null +++ b/src/models/starvla/oft_image_preprocess.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include + +namespace robotcpp::starvla { + +bool resize_pillow_bicubic_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, std::string & error); + +bool resize_torchvision_bicubic_aa_rgb(const uint8_t * source, int source_width, int source_height, int source_stride, + int target_width, int target_height, std::vector & target, + std::string & error); + +bool qwen3vl_smart_resize_dimensions(int source_width, int source_height, int factor, int min_pixels, int max_pixels, + int & target_width, int & target_height, std::string & error); + +bool preprocess_qwen3vl_rgb(const uint8_t * source, int source_width, int source_height, int channels, + int source_stride, int patch_size, int spatial_merge_size, int min_pixels, int max_pixels, + std::vector & target, int & target_width, int & target_height, + int & image_token_count, std::string & error); + +bool preprocess_oft_rgb(const uint8_t * source, int source_width, int source_height, int channels, int source_stride, + int training_width, int training_height, int processor_width, int processor_height, + std::vector & target, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.cpp b/src/models/starvla/oft_policy.cpp new file mode 100644 index 0000000..3f8bccc --- /dev/null +++ b/src/models/starvla/oft_policy.cpp @@ -0,0 +1,380 @@ +#include "models/starvla/oft_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" + +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +struct OFTBlockWeights { + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; + ggml_tensor * linear_weight = nullptr; + ggml_tensor * linear_bias = nullptr; +}; + +struct OFTWeights { + ggml_tensor * input_norm_weight = nullptr; + ggml_tensor * input_norm_bias = nullptr; + ggml_tensor * input_proj_weight = nullptr; + ggml_tensor * input_proj_bias = nullptr; + std::vector blocks; + ggml_tensor * output_norm_weight = nullptr; + ggml_tensor * output_norm_bias = nullptr; + ggml_tensor * output_proj_weight = nullptr; + ggml_tensor * output_proj_bias = nullptr; +}; + +using detail::has_shape; +using detail::require_bool; +using detail::require_f32; +using detail::require_i32; +using detail::require_string; +using detail::require_string_array; + +class OFTGGUFLoader final : public gguf_loader { + public: + OFTGGUFLoader(OFTPolicyConfig & config, OFTWeights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + const std::string architecture = require_string(gguf, "general.architecture"); + if (architecture != "starvla-policy") { + throw std::runtime_error("StarVLA policy GGUF has incompatible general.architecture: " + architecture); + } + if (require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "oft") { + throw std::runtime_error("StarVLA policy GGUF is not a supported Qwen OFT schema"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl" && config_.backbone_arch != "qwen2_5_vl") { + throw std::runtime_error("StarVLA OFT policy has an unsupported Qwen backbone"); + } + + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + if (config_.bundle_uuid.empty()) { + throw std::runtime_error("StarVLA policy bundle UUID is missing"); + } + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.text_filename.empty() || config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA policy component filenames must be non-empty"); + } + config_.input_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.hidden_dim = require_i32(gguf, "starvla.oft.hidden_size"); + config_.block_count = require_i32(gguf, "starvla.oft.block_count"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.layer_norm_epsilon = require_f32(gguf, "starvla.oft.layer_norm_epsilon"); + if (config_.input_dim <= 0 || config_.input_embedding_dim <= 0 || config_.vocab_size <= 0 || + config_.hidden_dim <= 0 || config_.block_count <= 0 || config_.action_dim <= 0 || config_.horizon <= 0 || + !std::isfinite(config_.layer_norm_epsilon) || config_.layer_norm_epsilon <= 0.0f) { + throw std::runtime_error("StarVLA OFT policy metadata has incompatible dimensions"); + } + + config_.prompt.horizon = config_.horizon; + config_.prompt.action_token = require_string(gguf, "starvla.prompt.action_token"); + config_.prompt.action_suffix = require_string(gguf, "starvla.prompt.action_suffix"); + config_.prompt.cot_enabled = require_bool(gguf, "starvla.prompt.cot_enabled"); + config_.prompt.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + config_.prompt.state_bins = require_i32(gguf, "starvla.prompt.state_bins"); + config_.prompt.state_bin_min = require_f32(gguf, "starvla.prompt.state_bin_min"); + config_.prompt.state_bin_max = require_f32(gguf, "starvla.prompt.state_bin_max"); + config_.prompt.state_clip = require_bool(gguf, "starvla.prompt.state_clip"); + config_.action_token_id = require_i32(gguf, "starvla.prompt.action_token_id"); + std::string prompt_error; + if (!validate_oft_prompt_config(config_.prompt, prompt_error) || config_.action_token_id < 0) { + throw std::runtime_error(prompt_error.empty() ? "StarVLA OFT token/template metadata is incompatible" + : prompt_error); + } + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + if (config_.image_count <= 0 || config_.image_names.size() != static_cast(config_.image_count) || + config_.image_processor_min_pixels <= 0 || + config_.image_processor_max_pixels < config_.image_processor_min_pixels || config_.image_patch_size <= 0 || + config_.image_spatial_merge_size <= 0 || config_.image_min_token_count <= 0 || + config_.image_max_token_count < config_.image_min_token_count) { + throw std::runtime_error("StarVLA OFT image metadata is incompatible"); + } + + config_.normalization = detail::require_normalization(gguf, config_.action_dim); + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + weights_.input_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.input_norm.weight"); + weights_.input_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.input_norm.bias"); + weights_.input_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.input_proj.weight"); + weights_.input_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.input_proj.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.oft.block." + std::to_string(block) + "."; + OFTBlockWeights current; + current.norm_weight = require_tensor(ctx_data, prefix + "norm.weight"); + current.norm_bias = require_tensor(ctx_data, prefix + "norm.bias"); + current.linear_weight = require_tensor(ctx_data, prefix + "linear.weight"); + current.linear_bias = require_tensor(ctx_data, prefix + "linear.bias"); + weights_.blocks.push_back(current); + } + weights_.output_norm_weight = require_tensor(ctx_data, "starvla.policy.oft.output_norm.weight"); + weights_.output_norm_bias = require_tensor(ctx_data, "starvla.policy.oft.output_norm.bias"); + weights_.output_proj_weight = require_tensor(ctx_data, "starvla.policy.oft.output_proj.weight"); + weights_.output_proj_bias = require_tensor(ctx_data, "starvla.policy.oft.output_proj.bias"); + + if (!has_shape(weights_.input_norm_weight, {config_.input_dim}) || + !has_shape(weights_.input_norm_bias, {config_.input_dim}) || + !has_shape(weights_.input_proj_weight, {config_.input_dim, config_.hidden_dim}) || + !has_shape(weights_.input_proj_bias, {config_.hidden_dim}) || + !has_shape(weights_.output_norm_weight, {config_.hidden_dim}) || + !has_shape(weights_.output_norm_bias, {config_.hidden_dim}) || + !has_shape(weights_.output_proj_weight, {config_.hidden_dim, config_.action_dim}) || + !has_shape(weights_.output_proj_bias, {config_.action_dim})) { + throw std::runtime_error("StarVLA OFT projection tensor has an incompatible ggml shape"); + } + for (const OFTBlockWeights & block : weights_.blocks) { + if (!has_shape(block.norm_weight, {config_.hidden_dim}) || + !has_shape(block.norm_bias, {config_.hidden_dim}) || + !has_shape(block.linear_weight, {config_.hidden_dim, config_.hidden_dim}) || + !has_shape(block.linear_bias, {config_.hidden_dim})) { + throw std::runtime_error("StarVLA OFT residual block tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + OFTPolicyConfig & config_; + OFTWeights & weights_; +}; + +} // namespace + +struct OFTPolicy::Impl { + OFTPolicyConfig config; + OFTWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + + ~Impl() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void build_graph() { + const size_t graph_size = GGML_DEFAULT_GRAPH_SIZE; + ggml_init_params params{}; + params.mem_size = graph_size * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA OFT graph context"); + } + + auto f32_vector = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto layer_norm = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * normalized = ggml_norm(graph_context, value, config.layer_norm_epsilon); + normalized = ggml_mul(graph_context, normalized, f32_vector(weight)); + return ggml_add(graph_context, normalized, f32_vector(bias)); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32_vector(bias)); + }; + + input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.input_dim, config.horizon); + ggml_set_name(input, "starvla_oft_action_queries"); + ggml_set_input(input); + + ggml_tensor * current = layer_norm(input, weights.input_norm_weight, weights.input_norm_bias); + current = ggml_relu(graph_context, linear(current, weights.input_proj_weight, weights.input_proj_bias)); + for (const OFTBlockWeights & block : weights.blocks) { + ggml_tensor * residual = current; + current = layer_norm(current, block.norm_weight, block.norm_bias); + current = ggml_relu(graph_context, linear(current, block.linear_weight, block.linear_bias)); + current = ggml_add(graph_context, current, residual); + } + current = layer_norm(current, weights.output_norm_weight, weights.output_norm_bias); + output = linear(current, weights.output_proj_weight, weights.output_proj_bias); + ggml_set_name(output, "starvla_oft_normalized_actions"); + ggml_set_output(output); + + graph = ggml_new_graph_custom(graph_context, graph_size, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA OFT graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA OFT graph"); + } + } +}; + +OFTPolicy::OFTPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +OFTPolicy::~OFTPolicy() = default; + +std::unique_ptr OFTPolicy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA OFT policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = GGML_DEFAULT_GRAPH_SIZE; + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { + error = "failed to initialize StarVLA OFT backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + OFTGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA OFT policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->build_graph(); + if (verbosity >= 1) { + std::fprintf(stderr, "%s: backend=%s input=%d hidden=%d blocks=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.input_dim, impl->config.hidden_dim, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new OFTPolicy(std::move(impl))); +} + +bool OFTPolicy::evaluate(const float * action_queries, size_t element_count, std::vector & normalized_actions, + std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr) { + error = "StarVLA OFT policy is not initialized"; + return false; + } + const size_t expected = static_cast(impl_->config.horizon) * static_cast(impl_->config.input_dim); + if (action_queries == nullptr || element_count != expected) { + error = "StarVLA OFT action-query tensor has an incompatible shape"; + return false; + } + + if (impl_->scheduler == nullptr || impl_->graph == nullptr || impl_->input == nullptr || impl_->output == nullptr) { + error = "StarVLA OFT policy graph is not initialized"; + return false; + } + + ggml_backend_tensor_set(impl_->input, action_queries, 0, element_count * sizeof(float)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA OFT graph compute failed"; + return false; + } + + const size_t output_count = + static_cast(impl_->config.horizon) * static_cast(impl_->config.action_dim); + normalized_actions.resize(output_count); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, output_count * sizeof(float)); + return true; +} + +bool OFTPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA OFT policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key, normalized_actions, impl_->config.horizon, + impl_->config.action_dim, actions, error); +} + +const OFTPolicyConfig & OFTPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA OFT policy is not initialized"); + } + return impl_->config; +} + +const char * OFTPolicy::backend_name() const { + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_policy.h b/src/models/starvla/oft_policy.h new file mode 100644 index 0000000..8416e1c --- /dev/null +++ b/src/models/starvla/oft_policy.h @@ -0,0 +1,64 @@ +#pragma once + +#include "models/starvla/normalization.h" +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct OFTPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + int input_dim = 0; + int input_embedding_dim = 0; + int vocab_size = 0; + int hidden_dim = 0; + int block_count = 0; + int action_dim = 0; + int horizon = 0; + float layer_norm_epsilon = 0.0f; + OFTPromptConfig prompt; + int action_token_id = 0; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + NormalizationConfig normalization; +}; + +class OFTPolicy { + public: + ~OFTPolicy(); + + OFTPolicy(const OFTPolicy &) = delete; + OFTPolicy & operator=(const OFTPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, std::string & error); + + bool evaluate(const float * action_queries, size_t element_count, std::vector & normalized_actions, + std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const OFTPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit OFTPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_prompt.cpp b/src/models/starvla/oft_prompt.cpp new file mode 100644 index 0000000..a2fc1da --- /dev/null +++ b/src/models/starvla/oft_prompt.cpp @@ -0,0 +1,161 @@ +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { +namespace { + +constexpr const char * kInstructionPlaceholder = "{instruction}"; +constexpr const char * kMtmdMediaMarker = "<__media__>"; + +std::string repeat(const std::string & value, int count) { + std::string result; + result.reserve(value.size() * static_cast(count)); + for (int i = 0; i < count; ++i) { + result += value; + } + return result; +} + +void replace_all(std::string & value, const std::string & needle, const std::string & replacement) { + size_t offset = 0; + while ((offset = value.find(needle, offset)) != std::string::npos) { + value.replace(offset, needle.size(), replacement); + offset += replacement.size(); + } +} + +bool discretize_state(const OFTPromptConfig & config, const std::vector & state, std::string & output, + std::string & error) { + if (state.empty()) { + output.clear(); + return true; + } + + std::ostringstream stream; + const double minimum = static_cast(config.state_bin_min); + const double maximum = static_cast(config.state_bin_max); + const double step = (maximum - minimum) / static_cast(config.state_bins); + for (size_t i = 0; i < state.size(); ++i) { + double value = static_cast(state[i]); + if (!std::isfinite(value)) { + error = "StarVLA OFT state contains a non-finite value"; + return false; + } + if (config.state_clip) { + value = std::max(minimum, std::min(maximum, value)); + } + + // Matches numpy.digitize(value, linspace(min, max, bins + 1)[:-1]) - 1. + int bin = -1; + for (int edge = 0; edge < config.state_bins; ++edge) { + const double boundary = minimum + step * static_cast(edge); + if (value >= boundary) { + bin = edge; + } else { + break; + } + } + if (i != 0) { + stream << ' '; + } + stream << bin; + } + output = stream.str(); + return true; +} + +} // namespace + +bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & error) { + error.clear(); + if (config.horizon <= 0 || config.action_token.empty()) { + error = "StarVLA OFT prompt has an invalid horizon or action token"; + return false; + } + const std::string expected_suffix = " Please predict the next " + std::to_string(config.horizon) + + " robot actions: " + repeat(config.action_token, config.horizon) + + "."; + if (config.action_suffix != expected_suffix) { + error = "StarVLA OFT action suffix does not match its horizon/token contract"; + return false; + } + if (config.cot_enabled && config.cot_template.find(kInstructionPlaceholder) == std::string::npos) { + error = "StarVLA OFT CoT template is missing {instruction}"; + return false; + } + if (config.state_bins <= 0 || !std::isfinite(config.state_bin_min) || !std::isfinite(config.state_bin_max) || + config.state_bin_max <= config.state_bin_min) { + error = "StarVLA OFT state prompt metadata is incompatible"; + return false; + } + return true; +} + +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, const std::vector & state, + std::string & instruction, std::string & error) { + instruction.clear(); + if (!validate_oft_prompt_config(config, error)) { + return false; + } + if (task.find(kMtmdMediaMarker) != std::string::npos) { + error = "StarVLA OFT task contains the reserved mtmd media marker"; + return false; + } + + instruction = task; + if (!state.empty()) { + std::string state_text; + if (!discretize_state(config, state, state_text, error)) { + instruction.clear(); + return false; + } + instruction += " [STATE] " + state_text + " [ACTION]"; + } + instruction += config.action_suffix; + + if (config.cot_enabled) { + std::string wrapped = config.cot_template; + replace_all(wrapped, kInstructionPlaceholder, instruction); + instruction = std::move(wrapped); + } + return true; +} + +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, const char * media_marker) { + const std::string marker = media_marker == nullptr ? std::string() : std::string(media_marker); + std::string content; + content.reserve(marker.size() * image_count + instruction.size()); + for (size_t i = 0; i < image_count; ++i) { + content += marker; + } + content += instruction; + return content; +} + +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, size_t count, + std::vector & positions, std::string & error) { + positions.clear(); + error.clear(); + if (count == 0) { + error = "StarVLA OFT action-token count must be positive"; + return false; + } + for (size_t i = 0; i < token_ids.size(); ++i) { + if (token_ids[i] == token_id) { + positions.push_back(i); + } + } + if (positions.size() < count) { + error = "StarVLA OFT prompt contains fewer action tokens than its horizon"; + positions.clear(); + return false; + } + positions.erase(positions.begin(), positions.end() - static_cast(count)); + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/oft_prompt.h b/src/models/starvla/oft_prompt.h new file mode 100644 index 0000000..393a7ea --- /dev/null +++ b/src/models/starvla/oft_prompt.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct OFTPromptConfig { + int horizon = 0; + std::string action_token; + std::string action_suffix; + bool cot_enabled = false; + std::string cot_template; + int state_bins = 0; + float state_bin_min = 0.0f; + float state_bin_max = 0.0f; + bool state_clip = false; +}; + +bool validate_oft_prompt_config(const OFTPromptConfig & config, std::string & error); + +bool build_oft_instruction(const OFTPromptConfig & config, const std::string & task, const std::vector & state, + std::string & instruction, std::string & error); + +std::string build_qwen_media_content(size_t image_count, const std::string & instruction, const char * media_marker); + +bool find_last_token_positions(const std::vector & token_ids, int32_t token_id, size_t count, + std::vector & positions, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_policy.cpp b/src/models/starvla/pi_policy.cpp new file mode 100644 index 0000000..70e24c7 --- /dev/null +++ b/src/models/starvla/pi_policy.cpp @@ -0,0 +1,699 @@ +#include "models/starvla/pi_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 16384; +struct PIBlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct PIWeights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + ggml_tensor * state_input_weight = nullptr; + ggml_tensor * state_input_bias = nullptr; + ggml_tensor * state_output_weight = nullptr; + ggml_tensor * state_output_bias = nullptr; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_i32_array; +using detail::require_string; +using detail::require_string_array; + +std::vector expected_hidden_tuple_indices(int qwen_layer_count, int block_count) { + std::vector result; + result.reserve(static_cast(block_count)); + const int first = qwen_layer_count + 1 - block_count; + for (int index = first; index <= qwen_layer_count; ++index) { + result.push_back(index); + } + return result; +} + +class PIGGUFLoader final : public gguf_loader { + public: + PIGGUFLoader(PIPolicyConfig & config, PIWeights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy" || + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "pi") { + throw std::runtime_error("GGUF is not a supported StarVLA PI policy"); + } + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.backbone_arch != "qwen2_5_vl" || config_.bundle_uuid.empty() || config_.text_filename.empty() || + config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA PI bundle metadata is incomplete"); + } + + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + config_.qwen_hidden_tuple_indices = require_i32_array(gguf, "starvla.conditioning.hidden_tuple_indices"); + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_framework_inference_pre_resize_width = + require_i32(gguf, "starvla.image.framework_inference_pre_resize_width"); + config_.image_framework_inference_pre_resize_height = + require_i32(gguf, "starvla.image.framework_inference_pre_resize_height"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi.block_count"); + config_.attention_head_count = require_i32(gguf, "starvla.pi.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.pi.attention_head_dim"); + config_.cross_attention_dim = require_i32(gguf, "starvla.pi.cross_attention_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.pi.feed_forward_dim"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.pi.mlp_hidden_dimension"); + config_.state_dim = require_i32(gguf, "starvla.state.dimension"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + config_.state_token_count = require_i32(gguf, "starvla.pi.state_token_count"); + config_.future_token_count = require_i32(gguf, "starvla.pi.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.pi.action_position_count"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.pi.timestep_projection_dim"); + config_.num_inference_timesteps = require_i32(gguf, "starvla.pi.num_inference_timesteps"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi.ada_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi.euler_dt"); + config_.timestep_ids = require_i32_array(gguf, "starvla.pi.timestep_ids"); + + const std::vector expected_indices = + expected_hidden_tuple_indices(config_.qwen_layer_count, config_.block_count); + const bool dimensions_valid = + config_.qwen_hidden_dim > 0 && config_.qwen_input_embedding_dim == config_.qwen_hidden_dim && + config_.qwen_layer_count >= config_.block_count && config_.qwen_vocab_size > 0 && config_.dit_width > 0 && + config_.dit_width % 2 == 0 && config_.block_count > 0 && config_.attention_head_count > 0 && + config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.cross_attention_dim == config_.qwen_hidden_dim && config_.feed_forward_dim > 0 && + config_.mlp_hidden_dim > 0 && config_.state_dim > 0 && config_.action_dim > 0 && config_.horizon > 0 && + config_.state_token_count == 1 && config_.future_token_count > 0 && + config_.action_position_count >= config_.horizon && config_.timestep_projection_dim >= 4 && + config_.timestep_projection_dim % 2 == 0 && config_.num_inference_timesteps > 0 && + config_.timestep_ids.size() == static_cast(config_.num_inference_timesteps) && + std::isfinite(config_.ada_norm_epsilon) && config_.ada_norm_epsilon > 0.0f && + std::isfinite(config_.euler_dt) && config_.euler_dt > 0.0f && + config_.qwen_hidden_tuple_indices == expected_indices && config_.image_count > 0 && + config_.image_names.size() == static_cast(config_.image_count) && + config_.image_framework_inference_pre_resize_width > 0 && + config_.image_framework_inference_pre_resize_height > 0 && config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count && !config_.cot_template.empty(); + if (!dimensions_valid) { + throw std::runtime_error("StarVLA PI dimensions, hidden taps, or sampler schedule are incompatible"); + } + + NormalizationConfig & normalization = config_.normalization; + normalization = detail::require_normalization(gguf, config_.action_dim); + if (!normalization.clip_actions || normalization.binary_comparison != "ge") { + throw std::runtime_error("StarVLA PI normalization must clip actions and use " + "binary comparison 'ge'"); + } + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, "starvla.policy.pi.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.pi.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.pi.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.pi.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.pi.block." + std::to_string(block) + "."; + PIBlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + bind(weights_.state_input_weight, "starvla.policy.pi.state.input.weight"); + bind(weights_.state_input_bias, "starvla.policy.pi.state.input.bias"); + bind(weights_.state_output_weight, "starvla.policy.pi.state.output.weight"); + bind(weights_.state_output_bias, "starvla.policy.pi.state.output.bias"); + bind(weights_.action_input_weight, "starvla.policy.pi.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.pi.action.input.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.pi.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.pi.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.pi.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.pi.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.pi.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.pi.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.pi.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.pi.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.pi.future_tokens.weight"); + bind(weights_.action_position, "starvla.policy.pi.action_position.weight"); + + const int64_t width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.state_input_weight, {config_.state_dim, config_.mlp_hidden_dim}) || + !has_shape(weights_.state_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.state_output_weight, {config_.mlp_hidden_dim, width}) || + !has_shape(weights_.state_output_bias, {width}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, {width, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA PI non-transformer tensor has an incompatible ggml shape"); + } + for (const PIBlockWeights & block : weights_.blocks) { + if (!has_shape(block.ada_norm_weight, {width, 2 * width}) || !has_shape(block.ada_norm_bias, {2 * width}) || + !has_shape(block.query_weight, {width, width}) || !has_shape(block.query_bias, {width}) || + !has_shape(block.key_weight, {config_.cross_attention_dim, width}) || + !has_shape(block.key_bias, {width}) || + !has_shape(block.value_weight, {config_.cross_attention_dim, width}) || + !has_shape(block.value_bias, {width}) || !has_shape(block.attention_output_weight, {width, width}) || + !has_shape(block.attention_output_bias, {width}) || + !has_shape(block.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(block.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(block.feed_forward_output_weight, {config_.feed_forward_dim, width}) || + !has_shape(block.feed_forward_output_bias, {width})) { + throw std::runtime_error("StarVLA PI transformer tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + PIPolicyConfig & config_; + PIWeights & weights_; +}; + +std::vector timestep_projection_table(const PIPolicyConfig & config) { + std::vector result(static_cast(config.num_inference_timesteps) * config.timestep_projection_dim); + const int half = config.timestep_projection_dim / 2; + for (int step = 0; step < config.num_inference_timesteps; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + for (int i = 0; i < half; ++i) { + const float exponent = -std::log(10000.0f) * i / static_cast(half - 1); + const float angle = timestep * std::exp(exponent); + const size_t offset = static_cast(step) * config.timestep_projection_dim; + result[offset + static_cast(i)] = std::cos(angle); + result[offset + static_cast(i + half)] = std::sin(angle); + } + } + return result; +} + +std::vector action_time_table(const PIPolicyConfig & config) { + std::vector result(static_cast(config.num_inference_timesteps) * config.dit_width); + const int half = config.dit_width / 2; + for (int step = 0; step < config.num_inference_timesteps; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + for (int i = 0; i < half; ++i) { + const float exponent = -std::log(10000.0f) * i / static_cast(half); + const float angle = timestep * std::exp(exponent); + const size_t offset = static_cast(step) * config.dit_width; + result[offset + static_cast(i)] = std::sin(angle); + result[offset + static_cast(i + half)] = std::cos(angle); + } + } + return result; +} + +} // namespace + +struct PIPolicy::Impl { + PIPolicyConfig config; + PIWeights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * state_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + bool graph_uses_state = false; + std::vector timestep_table; + std::vector action_table; + + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + state_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + graph_uses_state = false; + } + + void build_graph(size_t token_count, bool include_state) { + clear_graph(); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA PI conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA PI graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + hidden_input = ggml_new_tensor_3d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count), config.block_count); + if (include_state) { + state_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, config.state_dim); + } + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, + config.num_inference_timesteps); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, config.num_inference_timesteps); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || (include_state && state_input == nullptr) || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA PI graph inputs"); + } + ggml_set_name(hidden_input, "starvla_pi_qwen_hidden_states"); + if (state_input != nullptr) { + ggml_set_name(state_input, "starvla_pi_state"); + } + ggml_set_name(noise_input, "starvla_pi_initial_noise"); + ggml_set_name(timestep_projection_input, "starvla_pi_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_pi_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_pi_scalar_one"); + ggml_set_input(hidden_input); + if (state_input != nullptr) { + ggml_set_input(state_input); + } + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const PIBlockWeights & block) { + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + return ggml_add(graph_context, + ggml_mul(graph_context, normalized, ggml_add(graph_context, scale, scalar_one_input)), + shift); + }; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, const PIBlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = + ggml_soft_max_ext(graph_context, scores, nullptr, 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); + }; + + ggml_tensor * state_features = nullptr; + if (include_state) { + state_features = + ggml_relu(graph_context, linear(state_input, weights.state_input_weight, weights.state_input_bias)); + state_features = linear(state_features, weights.state_output_weight, weights.state_output_bias); + state_features = ggml_reshape_2d(graph_context, state_features, width, 1); + } + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + ggml_tensor * actions = noise_input; + + for (int step = 0; step < config.num_inference_timesteps; ++step) { + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); + + ggml_tensor * hidden = future; + if (state_features != nullptr) { + hidden = ggml_concat(graph_context, state_features, hidden, 1); + } + hidden = ggml_concat(graph_context, hidden, action_features, 1); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const PIBlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * layer_hidden = + ggml_view_2d(graph_context, hidden_input, config.qwen_hidden_dim, static_cast(token_count), + hidden_input->nb[1], static_cast(block_index) * hidden_input->nb[2]); + ggml_tensor * normalized = ada_norm(hidden, temb, block); + hidden = ggml_add(graph_context, hidden, attention(normalized, layer_hidden, block)); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + hidden = + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); + ggml_tensor * velocity = ggml_view_2d( + graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast((include_state ? config.state_token_count : 0) + config.future_token_count) * + hidden->nb[1]); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_pi_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA PI graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA PI graph"); + } + conditioning_token_count = token_count; + graph_uses_state = include_state; + } +}; + +PIPolicy::PIPolicy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +PIPolicy::~PIPolicy() = default; + +std::unique_ptr PIPolicy::load(const std::string & path, int n_threads, int verbosity, std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA PI policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { + error = "failed to initialize StarVLA PI backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + PIGGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA PI policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf(stderr, + "%s: backend=%s qwen=%d width=%d blocks=%d horizon=%d " + "action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new PIPolicy(std::move(impl))); +} + +bool PIPolicy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const float * state, + size_t state_element_count, const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA PI policy is not initialized"; + return false; + } + const size_t layer_width = static_cast(impl_->config.block_count) * impl_->config.qwen_hidden_dim; + if (qwen_hidden_states == nullptr || layer_width == 0 || hidden_element_count == 0 || + hidden_element_count % layer_width != 0) { + error = "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + return false; + } + const size_t token_count = hidden_element_count / layer_width; + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + error = "StarVLA PI layer-wise Qwen conditioning tensor has an incompatible shape"; + return false; + } + const bool include_state = state_element_count != 0; + if (include_state && (state == nullptr || state_element_count != static_cast(impl_->config.state_dim))) { + error = "StarVLA PI state tensor has an incompatible shape"; + return false; + } + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (initial_noise == nullptr || noise_element_count != expected_noise) { + error = "StarVLA PI initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + (include_state && + std::any_of(state, state + state_element_count, [](float value) { return !std::isfinite(value); })) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = "StarVLA PI conditioning, state, and initial noise must be finite"; + return false; + } + + try { + if (impl_->graph == nullptr || impl_->conditioning_token_count != token_count || + impl_->graph_uses_state != include_state) { + impl_->build_graph(token_count, include_state); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); + if (include_state) { + ggml_backend_tensor_set(impl_->state_input, state, 0, state_element_count * sizeof(float)); + } + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA PI graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA PI graph produced non-finite actions"; + return false; + } + return true; +} + +bool PIPolicy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA PI policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); +} + +const PIPolicyConfig & PIPolicy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA PI policy is not initialized"); + } + return impl_->config; +} + +const char * PIPolicy::backend_name() const { + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_policy.h b/src/models/starvla/pi_policy.h new file mode 100644 index 0000000..0b7b2a4 --- /dev/null +++ b/src/models/starvla/pi_policy.h @@ -0,0 +1,89 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct PIPolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; + std::string cot_template; + std::vector qwen_hidden_tuple_indices; + int image_count = 0; + std::vector image_names; + int image_framework_inference_pre_resize_width = 0; + int image_framework_inference_pre_resize_height = 0; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + int dit_width = 0; + int block_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int cross_attention_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int state_dim = 0; + int action_dim = 0; + int horizon = 0; + int state_token_count = 0; + int future_token_count = 0; + int action_position_count = 0; + int timestep_projection_dim = 0; + int num_inference_timesteps = 0; + float ada_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class PIPolicy { + public: + ~PIPolicy(); + + PIPolicy(const PIPolicy &) = delete; + PIPolicy & operator=(const PIPolicy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, std::string & error); + + // qwen_hidden_states is layer-major + // [block_count, token_count, qwen_hidden_dim]. The legacy released + // implementation did not forward the Qwen attention mask into the policy + // head. state is either omitted (the official Bridge deployment path) or + // one token [state_dim], and initial_noise is token-major + // [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const float * state, + size_t state_element_count, const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const PIPolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit PIPolicy(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_v3_policy.cpp b/src/models/starvla/pi_v3_policy.cpp new file mode 100644 index 0000000..bedc679 --- /dev/null +++ b/src/models/starvla/pi_v3_policy.cpp @@ -0,0 +1,744 @@ +#include "models/starvla/pi_v3_policy.h" + +#include "ggml-backend.h" +#include "ggml.h" +#include "gguf.h" +#include "models/ggml_backend.h" +#include "models/gguf_loader.h" +#include "models/starvla/policy_gguf.h" +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +namespace { + +constexpr size_t kGraphSize = 32768; +constexpr int kKQMaskPad = 32; +constexpr int kReleasedLayerCount = 36; + +struct PIV3BlockWeights { + ggml_tensor * ada_norm_weight = nullptr; + ggml_tensor * ada_norm_bias = nullptr; + ggml_tensor * query_weight = nullptr; + ggml_tensor * query_bias = nullptr; + ggml_tensor * key_weight = nullptr; + ggml_tensor * key_bias = nullptr; + ggml_tensor * value_weight = nullptr; + ggml_tensor * value_bias = nullptr; + ggml_tensor * attention_output_weight = nullptr; + ggml_tensor * attention_output_bias = nullptr; + ggml_tensor * feed_forward_input_weight = nullptr; + ggml_tensor * feed_forward_input_bias = nullptr; + ggml_tensor * feed_forward_output_weight = nullptr; + ggml_tensor * feed_forward_output_bias = nullptr; +}; + +struct PIV3ProjectorWeights { + ggml_tensor * norm_weight = nullptr; + ggml_tensor * norm_bias = nullptr; + ggml_tensor * projection_weight = nullptr; + ggml_tensor * projection_bias = nullptr; +}; + +struct PIV3Weights { + ggml_tensor * timestep_input_weight = nullptr; + ggml_tensor * timestep_input_bias = nullptr; + ggml_tensor * timestep_output_weight = nullptr; + ggml_tensor * timestep_output_bias = nullptr; + std::vector blocks; + std::vector projectors; + ggml_tensor * action_input_weight = nullptr; + ggml_tensor * action_input_bias = nullptr; + ggml_tensor * action_time_mix_weight = nullptr; + ggml_tensor * action_time_mix_bias = nullptr; + ggml_tensor * action_output_weight = nullptr; + ggml_tensor * action_output_bias = nullptr; + ggml_tensor * velocity_input_weight = nullptr; + ggml_tensor * velocity_input_bias = nullptr; + ggml_tensor * velocity_output_weight = nullptr; + ggml_tensor * velocity_output_bias = nullptr; + ggml_tensor * future_tokens = nullptr; + ggml_tensor * action_position = nullptr; +}; + +using detail::has_shape; +using detail::require_f32; +using detail::require_i32; +using detail::require_string; +using detail::require_string_array; + +std::vector integer_range(int first, int count) { + std::vector result(static_cast(count)); + for (int index = 0; index < count; ++index) { + result[static_cast(index)] = first + index; + } + return result; +} + +class PIV3GGUFLoader final : public gguf_loader { + public: + PIV3GGUFLoader(PIV3PolicyConfig & config, PIV3Weights & weights) : config_(config), weights_(weights) {} + + protected: + bool parse_metadata(gguf_context * gguf) override { + if (require_string(gguf, "general.architecture") != "starvla-policy" || + require_i32(gguf, "starvla.schema_version") != 1 || require_string(gguf, "starvla.framework") != "pi_v3") { + throw std::runtime_error("GGUF is not a supported StarVLA PI-v3 policy"); + } + + config_.backbone_arch = require_string(gguf, "starvla.backbone.arch"); + if (config_.backbone_arch != "qwen3_vl") { + throw std::runtime_error("StarVLA PI-v3 requires a Qwen3-VL backbone"); + } + config_.bundle_uuid = require_string(gguf, "starvla.bundle.uuid"); + config_.text_filename = require_string(gguf, "starvla.component.text.filename"); + config_.mmproj_filename = require_string(gguf, "starvla.component.mmproj.filename"); + if (config_.bundle_uuid.empty() || config_.text_filename.empty() || config_.mmproj_filename.empty()) { + throw std::runtime_error("StarVLA PI-v3 bundle metadata is incomplete"); + } + + config_.qwen_hidden_dim = require_i32(gguf, "starvla.qwen.hidden_size"); + config_.qwen_input_embedding_dim = require_i32(gguf, "starvla.qwen.input_embedding_size"); + config_.qwen_layer_count = require_i32(gguf, "starvla.qwen.layer_count"); + config_.qwen_vocab_size = require_i32(gguf, "starvla.qwen.vocab_size"); + config_.cot_template = require_string(gguf, "starvla.prompt.cot_template"); + + config_.image_count = require_i32(gguf, "starvla.image.count"); + config_.image_names = require_string_array(gguf, "starvla.image.names"); + config_.image_processor_min_pixels = require_i32(gguf, "starvla.image.processor_min_pixels"); + config_.image_processor_max_pixels = require_i32(gguf, "starvla.image.processor_max_pixels"); + config_.image_patch_size = require_i32(gguf, "starvla.image.patch_size"); + config_.image_spatial_merge_size = require_i32(gguf, "starvla.image.spatial_merge_size"); + config_.image_min_token_count = require_i32(gguf, "starvla.image.min_token_count"); + config_.image_max_token_count = require_i32(gguf, "starvla.image.max_token_count"); + + config_.dit_width = require_i32(gguf, "starvla.pi_v3.dit_width"); + config_.block_count = require_i32(gguf, "starvla.pi_v3.block_count"); + config_.projector_count = require_i32(gguf, "starvla.pi_v3.projector_count"); + config_.attention_head_count = require_i32(gguf, "starvla.pi_v3.attention_head_count"); + config_.attention_head_dim = require_i32(gguf, "starvla.pi_v3.attention_head_dim"); + config_.feed_forward_dim = require_i32(gguf, "starvla.pi_v3.feed_forward_dim"); + config_.mlp_hidden_dim = require_i32(gguf, "starvla.pi_v3.mlp_hidden_dimension"); + config_.future_token_count = require_i32(gguf, "starvla.pi_v3.future_token_count"); + config_.action_position_count = require_i32(gguf, "starvla.pi_v3.action_position_count"); + config_.no_state_sequence_length = require_i32(gguf, "starvla.pi_v3.no_state_sequence_length"); + config_.timestep_projection_dim = require_i32(gguf, "starvla.pi_v3.timestep_projection_dim"); + config_.num_timestep_buckets = require_i32(gguf, "starvla.pi_v3.num_timestep_buckets"); + config_.num_inference_timesteps = require_i32(gguf, "starvla.pi_v3.num_inference_timesteps"); + config_.ada_norm_epsilon = require_f32(gguf, "starvla.pi_v3.ada_norm_epsilon"); + config_.projector_norm_epsilon = require_f32(gguf, "starvla.pi_v3.projector_norm_epsilon"); + config_.euler_dt = require_f32(gguf, "starvla.pi_v3.euler_dt"); + config_.action_dim = require_i32(gguf, "starvla.action.dimension"); + config_.horizon = require_i32(gguf, "starvla.action.horizon"); + + const bool valid = + config_.qwen_hidden_dim > 0 && config_.qwen_input_embedding_dim > 0 && + config_.qwen_layer_count == kReleasedLayerCount && config_.qwen_vocab_size > 0 && + !config_.cot_template.empty() && config_.image_count > 0 && + config_.image_names.size() == static_cast(config_.image_count) && + config_.image_processor_min_pixels > 0 && + config_.image_processor_max_pixels >= config_.image_processor_min_pixels && config_.image_patch_size > 0 && + config_.image_spatial_merge_size > 0 && config_.image_min_token_count > 0 && + config_.image_max_token_count >= config_.image_min_token_count && config_.dit_width > 0 && + config_.block_count == kReleasedLayerCount && config_.projector_count == config_.block_count && + config_.attention_head_count > 0 && config_.attention_head_dim > 0 && + config_.attention_head_count * config_.attention_head_dim == config_.dit_width && + config_.feed_forward_dim > 0 && config_.mlp_hidden_dim > 0 && config_.action_dim > 0 && + config_.horizon > 0 && config_.future_token_count > 0 && config_.action_position_count >= config_.horizon && + config_.no_state_sequence_length == config_.future_token_count + config_.horizon && + config_.timestep_projection_dim >= 4 && config_.timestep_projection_dim % 2 == 0 && + config_.num_timestep_buckets > 0 && config_.num_inference_timesteps == 4 && + config_.ada_norm_epsilon > 0.0f && config_.projector_norm_epsilon > 0.0f && config_.euler_dt > 0.0f; + if (!valid) { + throw std::runtime_error("StarVLA PI-v3 metadata has incompatible dimensions"); + } + + config_.qwen_hidden_tuple_indices = integer_range(1, config_.qwen_layer_count); + config_.timestep_ids.resize(static_cast(config_.num_inference_timesteps)); + for (int step = 0; step < config_.num_inference_timesteps; ++step) { + config_.timestep_ids[static_cast(step)] = + step * config_.num_timestep_buckets / config_.num_inference_timesteps; + } + + config_.normalization = detail::require_normalization(gguf, config_.action_dim); + return true; + } + + bool bind_tensors(ggml_context * ctx_data) override { + auto bind = [&](ggml_tensor *& destination, const std::string & name) { + destination = require_tensor(ctx_data, name); + }; + bind(weights_.timestep_input_weight, "starvla.policy.pi_v3.timestep.input.weight"); + bind(weights_.timestep_input_bias, "starvla.policy.pi_v3.timestep.input.bias"); + bind(weights_.timestep_output_weight, "starvla.policy.pi_v3.timestep.output.weight"); + bind(weights_.timestep_output_bias, "starvla.policy.pi_v3.timestep.output.bias"); + weights_.blocks.clear(); + weights_.blocks.reserve(static_cast(config_.block_count)); + for (int block = 0; block < config_.block_count; ++block) { + const std::string prefix = "starvla.policy.pi_v3.block." + std::to_string(block) + "."; + PIV3BlockWeights current; + bind(current.ada_norm_weight, prefix + "ada_norm.weight"); + bind(current.ada_norm_bias, prefix + "ada_norm.bias"); + bind(current.query_weight, prefix + "attention.query.weight"); + bind(current.query_bias, prefix + "attention.query.bias"); + bind(current.key_weight, prefix + "attention.key.weight"); + bind(current.key_bias, prefix + "attention.key.bias"); + bind(current.value_weight, prefix + "attention.value.weight"); + bind(current.value_bias, prefix + "attention.value.bias"); + bind(current.attention_output_weight, prefix + "attention.output.weight"); + bind(current.attention_output_bias, prefix + "attention.output.bias"); + bind(current.feed_forward_input_weight, prefix + "feed_forward.input.weight"); + bind(current.feed_forward_input_bias, prefix + "feed_forward.input.bias"); + bind(current.feed_forward_output_weight, prefix + "feed_forward.output.weight"); + bind(current.feed_forward_output_bias, prefix + "feed_forward.output.bias"); + weights_.blocks.push_back(current); + } + weights_.projectors.clear(); + weights_.projectors.reserve(static_cast(config_.projector_count)); + for (int projector = 0; projector < config_.projector_count; ++projector) { + const std::string prefix = "starvla.policy.pi_v3.projector." + std::to_string(projector) + "."; + PIV3ProjectorWeights current; + bind(current.norm_weight, prefix + "norm.weight"); + bind(current.norm_bias, prefix + "norm.bias"); + bind(current.projection_weight, prefix + "projection.weight"); + bind(current.projection_bias, prefix + "projection.bias"); + weights_.projectors.push_back(current); + } + bind(weights_.action_input_weight, "starvla.policy.pi_v3.action.input.weight"); + bind(weights_.action_input_bias, "starvla.policy.pi_v3.action.input.bias"); + bind(weights_.action_time_mix_weight, "starvla.policy.pi_v3.action.time_mix.weight"); + bind(weights_.action_time_mix_bias, "starvla.policy.pi_v3.action.time_mix.bias"); + bind(weights_.action_output_weight, "starvla.policy.pi_v3.action.output.weight"); + bind(weights_.action_output_bias, "starvla.policy.pi_v3.action.output.bias"); + bind(weights_.velocity_input_weight, "starvla.policy.pi_v3.velocity.input.weight"); + bind(weights_.velocity_input_bias, "starvla.policy.pi_v3.velocity.input.bias"); + bind(weights_.velocity_output_weight, "starvla.policy.pi_v3.velocity.output.weight"); + bind(weights_.velocity_output_bias, "starvla.policy.pi_v3.velocity.output.bias"); + bind(weights_.future_tokens, "starvla.policy.pi_v3.future_tokens.weight"); + bind(weights_.action_position, "starvla.policy.pi_v3.action_position.weight"); + + const int width = config_.dit_width; + if (!has_shape(weights_.timestep_input_weight, {config_.timestep_projection_dim, width}) || + !has_shape(weights_.timestep_input_bias, {width}) || + !has_shape(weights_.timestep_output_weight, {width, width}) || + !has_shape(weights_.timestep_output_bias, {width}) || + !has_shape(weights_.action_input_weight, {config_.action_dim, width}) || + !has_shape(weights_.action_input_bias, {width}) || + !has_shape(weights_.action_time_mix_weight, {2 * width, width}) || + !has_shape(weights_.action_time_mix_bias, {width}) || + !has_shape(weights_.action_output_weight, {width, width}) || + !has_shape(weights_.action_output_bias, {width}) || + !has_shape(weights_.velocity_input_weight, {width, config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_input_bias, {config_.mlp_hidden_dim}) || + !has_shape(weights_.velocity_output_weight, {config_.mlp_hidden_dim, config_.action_dim}) || + !has_shape(weights_.velocity_output_bias, {config_.action_dim}) || + !has_shape(weights_.future_tokens, {width, config_.future_token_count}) || + !has_shape(weights_.action_position, {width, config_.action_position_count})) { + throw std::runtime_error("StarVLA PI_v3 non-block tensor has an incompatible ggml shape"); + } + for (const PIV3BlockWeights & current : weights_.blocks) { + if (!has_shape(current.ada_norm_weight, {width, 2 * width}) || + !has_shape(current.ada_norm_bias, {2 * width}) || !has_shape(current.query_weight, {width, width}) || + !has_shape(current.query_bias, {width}) || !has_shape(current.key_weight, {width, width}) || + !has_shape(current.key_bias, {width}) || !has_shape(current.value_weight, {width, width}) || + !has_shape(current.value_bias, {width}) || + !has_shape(current.attention_output_weight, {width, width}) || + !has_shape(current.attention_output_bias, {width}) || + !has_shape(current.feed_forward_input_weight, {width, config_.feed_forward_dim}) || + !has_shape(current.feed_forward_input_bias, {config_.feed_forward_dim}) || + !has_shape(current.feed_forward_output_weight, {config_.feed_forward_dim, width}) || + !has_shape(current.feed_forward_output_bias, {width})) { + throw std::runtime_error("StarVLA PI_v3 transformer block tensor has an incompatible ggml shape"); + } + } + for (const PIV3ProjectorWeights & current : weights_.projectors) { + if (!has_shape(current.norm_weight, {config_.qwen_hidden_dim}) || + !has_shape(current.norm_bias, {config_.qwen_hidden_dim}) || + !has_shape(current.projection_weight, {config_.qwen_hidden_dim, width}) || + !has_shape(current.projection_bias, {width})) { + throw std::runtime_error("StarVLA PI_v3 projector tensor has an incompatible ggml shape"); + } + } + return true; + } + + private: + PIV3PolicyConfig & config_; + PIV3Weights & weights_; +}; + +std::vector timestep_projection_table(const PIV3PolicyConfig & config) { + const int dim = config.timestep_projection_dim; + const int half = dim / 2; + const float denominator = static_cast(half - 1); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::cos(angle); + row[index + half] = std::sin(angle); + } + } + return table; +} + +std::vector action_time_table(const PIV3PolicyConfig & config) { + const int dim = config.dit_width; + const int half = dim / 2; + const float denominator = static_cast(half); + std::vector table(static_cast(dim) * 4, 0.0f); + for (int step = 0; step < 4; ++step) { + const float timestep = static_cast(config.timestep_ids[static_cast(step)]); + float * row = table.data() + static_cast(step) * dim; + for (int index = 0; index < half; ++index) { + const float frequency = std::exp(-std::log(10000.0f) * static_cast(index) / denominator); + const float angle = timestep * frequency; + row[index] = std::sin(angle); + row[index + half] = std::cos(angle); + } + } + return table; +} + +} // namespace + +struct PIV3Policy::Impl { + PIV3PolicyConfig config; + PIV3Weights weights; + gguf_load_result loaded; + ggml_backend_t backend_cpu = nullptr; + std::vector backends; + ggml_backend_sched_t scheduler = nullptr; + backend_buft_policy buft_policy; + backend_mode mode = backend_mode::cpu; + int n_threads = 0; + int verbosity = 0; + ggml_context * graph_context = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * hidden_input = nullptr; + ggml_tensor * cross_mask_input = nullptr; + ggml_tensor * noise_input = nullptr; + ggml_tensor * timestep_projection_input = nullptr; + ggml_tensor * action_time_input = nullptr; + ggml_tensor * scalar_one_input = nullptr; + ggml_tensor * output = nullptr; + size_t conditioning_token_count = 0; + std::vector timestep_table; + std::vector action_table; + ~Impl() { + clear_graph(); + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_free(scheduler); + scheduler = nullptr; + } + if (loaded.model_buffer != nullptr) { + ggml_backend_buffer_free(loaded.model_buffer); + loaded.model_buffer = nullptr; + } + if (loaded.ctx_data != nullptr) { + ggml_free(loaded.ctx_data); + loaded.ctx_data = nullptr; + } + if (loaded.gguf != nullptr) { + gguf_free(loaded.gguf); + loaded.gguf = nullptr; + } + for (ggml_backend_t backend : backends) { + if (backend != nullptr) { + ggml_backend_free(backend); + } + } + backends.clear(); + backend_cpu = nullptr; + } + + void clear_graph() { + if (scheduler != nullptr) { + ggml_backend_sched_synchronize(scheduler); + ggml_backend_sched_reset(scheduler); + } + if (graph_context != nullptr) { + ggml_free(graph_context); + graph_context = nullptr; + } + graph = nullptr; + hidden_input = nullptr; + cross_mask_input = nullptr; + noise_input = nullptr; + timestep_projection_input = nullptr; + action_time_input = nullptr; + scalar_one_input = nullptr; + output = nullptr; + conditioning_token_count = 0; + } + + void build_graph(size_t token_count) { + clear_graph(); + if (token_count == 0 || token_count > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("invalid StarVLA PI_v3 conditioning token count"); + } + + ggml_init_params params{}; + params.mem_size = kGraphSize * ggml_tensor_overhead() + ggml_graph_overhead_custom(kGraphSize, false); + params.mem_buffer = nullptr; + params.no_alloc = true; + graph_context = ggml_init(params); + if (graph_context == nullptr) { + throw std::runtime_error("failed to initialize StarVLA PI_v3 graph context"); + } + + const int width = config.dit_width; + const int heads = config.attention_head_count; + const int head_dim = config.attention_head_dim; + const int sequence_length = config.no_state_sequence_length; + const int mask_queries = GGML_PAD(sequence_length, kKQMaskPad); + + hidden_input = ggml_new_tensor_3d(graph_context, GGML_TYPE_F32, config.qwen_hidden_dim, + static_cast(token_count), config.qwen_layer_count); + cross_mask_input = + ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, static_cast(token_count), mask_queries); + noise_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.action_dim, config.horizon); + timestep_projection_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, config.timestep_projection_dim, 4); + action_time_input = ggml_new_tensor_2d(graph_context, GGML_TYPE_F32, width, 4); + scalar_one_input = ggml_new_tensor_1d(graph_context, GGML_TYPE_F32, 1); + if (hidden_input == nullptr || cross_mask_input == nullptr || noise_input == nullptr || + timestep_projection_input == nullptr || action_time_input == nullptr || scalar_one_input == nullptr) { + throw std::runtime_error("failed to create StarVLA PI_v3 graph inputs"); + } + ggml_set_name(hidden_input, "starvla_pi_v3_qwen_hidden_states"); + ggml_set_name(cross_mask_input, "starvla_pi_v3_qwen_attention_mask"); + ggml_set_name(noise_input, "starvla_pi_v3_initial_noise"); + ggml_set_name(timestep_projection_input, "starvla_pi_v3_timestep_projection_table"); + ggml_set_name(action_time_input, "starvla_pi_v3_action_time_table"); + ggml_set_name(scalar_one_input, "starvla_pi_v3_scalar_one"); + ggml_set_input(hidden_input); + ggml_set_input(cross_mask_input); + ggml_set_input(noise_input); + ggml_set_input(timestep_projection_input); + ggml_set_input(action_time_input); + ggml_set_input(scalar_one_input); + + auto f32 = [&](ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_F32 ? tensor : ggml_cast(graph_context, tensor, GGML_TYPE_F32); + }; + auto bf16_roundtrip = [&](ggml_tensor * tensor) { + return ggml_cast(graph_context, ggml_cast(graph_context, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); + }; + auto linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * projected = ggml_mul_mat(graph_context, weight, value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + return ggml_add(graph_context, projected, f32(bias)); + }; + auto projector_linear = [&](ggml_tensor * value, ggml_tensor * weight, ggml_tensor * bias) { + ggml_tensor * bf16_value = ggml_cast(graph_context, value, GGML_TYPE_BF16); + ggml_tensor * bf16_weight = ggml_cast(graph_context, weight, GGML_TYPE_BF16); + ggml_tensor * projected = ggml_mul_mat(graph_context, bf16_weight, bf16_value); + ggml_mul_mat_set_prec(projected, GGML_PREC_F32); + projected = ggml_add(graph_context, projected, bf16_roundtrip(bias)); + return bf16_roundtrip(projected); + }; + auto ada_norm = [&](ggml_tensor * value, ggml_tensor * temb, const PIV3BlockWeights & block) { + ggml_tensor * modulation = + linear(ggml_silu(graph_context, temb), block.ada_norm_weight, block.ada_norm_bias); + ggml_tensor * scale = ggml_view_1d(graph_context, modulation, width, 0); + ggml_tensor * shift = + ggml_view_1d(graph_context, modulation, width, static_cast(width) * sizeof(float)); + ggml_tensor * normalized = ggml_norm(graph_context, value, config.ada_norm_epsilon); + return ggml_add(graph_context, + ggml_mul(graph_context, normalized, ggml_add(graph_context, scale, scalar_one_input)), + shift); + }; + auto attention = [&](ggml_tensor * query_source, ggml_tensor * key_value_source, + const PIV3BlockWeights & block) { + const int64_t query_count = query_source->ne[1]; + const int64_t key_value_count = key_value_source->ne[1]; + ggml_tensor * query = linear(query_source, block.query_weight, block.query_bias); + ggml_tensor * key = linear(key_value_source, block.key_weight, block.key_bias); + ggml_tensor * value = linear(key_value_source, block.value_weight, block.value_bias); + query = ggml_reshape_3d(graph_context, query, head_dim, heads, query_count); + key = ggml_reshape_3d(graph_context, key, head_dim, heads, key_value_count); + value = ggml_reshape_3d(graph_context, value, head_dim, heads, key_value_count); + query = ggml_permute(graph_context, query, 0, 2, 1, 3); + key = ggml_permute(graph_context, key, 0, 2, 1, 3); + value = ggml_cont(graph_context, ggml_permute(graph_context, value, 1, 2, 0, 3)); + ggml_tensor * scores = ggml_mul_mat(graph_context, key, query); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_soft_max_ext(graph_context, scores, cross_mask_input, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * attended = ggml_mul_mat(graph_context, value, scores); + ggml_mul_mat_set_prec(attended, GGML_PREC_F32); + attended = ggml_permute(graph_context, attended, 0, 2, 1, 3); + attended = ggml_cont_2d(graph_context, attended, width, query_count); + return linear(attended, block.attention_output_weight, block.attention_output_bias); + }; + + std::vector projected_hidden_states; + projected_hidden_states.reserve(static_cast(config.projector_count)); + for (int layer = 0; layer < config.projector_count; ++layer) { + const PIV3ProjectorWeights & projector = weights.projectors[static_cast(layer)]; + ggml_tensor * layer_hidden = + ggml_view_2d(graph_context, hidden_input, config.qwen_hidden_dim, static_cast(token_count), + hidden_input->nb[1], static_cast(layer) * hidden_input->nb[2]); + layer_hidden = bf16_roundtrip(layer_hidden); + layer_hidden = ggml_norm(graph_context, layer_hidden, config.projector_norm_epsilon); + layer_hidden = ggml_mul(graph_context, layer_hidden, f32(projector.norm_weight)); + layer_hidden = ggml_add(graph_context, layer_hidden, f32(projector.norm_bias)); + ggml_tensor * projected = + projector_linear(layer_hidden, projector.projection_weight, projector.projection_bias); + projected_hidden_states.push_back(projected); + } + + ggml_tensor * future = f32(weights.future_tokens); + ggml_tensor * position_view = ggml_view_2d(graph_context, weights.action_position, width, config.horizon, + weights.action_position->nb[1], 0); + ggml_tensor * position = f32(position_view); + // Qwen/projector inference and torch.randn run at BF16 in the released + // script. The action head then enters CUDA autocast(float32). + ggml_tensor * actions = bf16_roundtrip(noise_input); + + for (int step = 0; step < 4; ++step) { + ggml_tensor * timestep_projection = + ggml_view_1d(graph_context, timestep_projection_input, config.timestep_projection_dim, + static_cast(step) * config.timestep_projection_dim * sizeof(float)); + ggml_tensor * temb = + linear(timestep_projection, weights.timestep_input_weight, weights.timestep_input_bias); + temb = ggml_silu(graph_context, temb); + temb = linear(temb, weights.timestep_output_weight, weights.timestep_output_bias); + + ggml_tensor * action_features = linear(actions, weights.action_input_weight, weights.action_input_bias); + ggml_tensor * action_time = ggml_view_1d(graph_context, action_time_input, width, + static_cast(step) * width * sizeof(float)); + action_time = ggml_repeat(graph_context, action_time, action_features); + action_features = ggml_concat(graph_context, action_features, action_time, 0); + action_features = linear(action_features, weights.action_time_mix_weight, weights.action_time_mix_bias); + action_features = ggml_silu(graph_context, action_features); + action_features = linear(action_features, weights.action_output_weight, weights.action_output_bias); + action_features = ggml_add(graph_context, action_features, position); + + ggml_tensor * hidden = ggml_concat(graph_context, future, action_features, 1); + for (int block_index = 0; block_index < config.block_count; ++block_index) { + const PIV3BlockWeights & block = weights.blocks[static_cast(block_index)]; + ggml_tensor * normalized = ada_norm(hidden, temb, block); + ggml_tensor * attended = + attention(normalized, projected_hidden_states[static_cast(block_index)], block); + hidden = ggml_add(graph_context, hidden, attended); + ggml_tensor * ff = ggml_norm(graph_context, hidden, config.ada_norm_epsilon); + ff = linear(ff, block.feed_forward_input_weight, block.feed_forward_input_bias); + ff = ggml_gelu(graph_context, ff); + ff = linear(ff, block.feed_forward_output_weight, block.feed_forward_output_bias); + hidden = ggml_add(graph_context, hidden, ff); + } + + // The released legacy sampler calls DiT with return_pre_output=true. + // norm_out/proj_out_1/proj_out_2 are therefore intentionally inactive. + hidden = + ggml_relu(graph_context, linear(hidden, weights.velocity_input_weight, weights.velocity_input_bias)); + hidden = linear(hidden, weights.velocity_output_weight, weights.velocity_output_bias); + ggml_tensor * velocity = + ggml_view_2d(graph_context, hidden, config.action_dim, config.horizon, hidden->nb[1], + static_cast(config.future_token_count) * hidden->nb[1]); + actions = ggml_add(graph_context, actions, ggml_scale(graph_context, velocity, config.euler_dt)); + } + + output = actions; + ggml_set_name(output, "starvla_pi_v3_normalized_actions"); + ggml_set_output(output); + graph = ggml_new_graph_custom(graph_context, kGraphSize, false); + if (graph == nullptr) { + throw std::runtime_error("failed to create StarVLA PI_v3 graph"); + } + ggml_build_forward_expand(graph, output); + ggml_backend_sched_reset(scheduler); + if (!ggml_backend_sched_alloc_graph(scheduler, graph)) { + throw std::runtime_error("failed to allocate StarVLA PI_v3 graph"); + } + + conditioning_token_count = token_count; + } +}; + +PIV3Policy::PIV3Policy(std::unique_ptr impl) : impl_(std::move(impl)) {} + +PIV3Policy::~PIV3Policy() = default; + +std::unique_ptr PIV3Policy::load(const std::string & path, int n_threads, int verbosity, + std::string & error) { + error.clear(); + if (path.empty()) { + error = "StarVLA PI_v3 policy path is required"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->n_threads = n_threads; + impl->verbosity = verbosity; + try { + backend_scheduler_config scheduler_config; + scheduler_config.max_nodes = static_cast(kGraphSize); + scheduler_config.parallel = false; + scheduler_config.op_offload = true; + backend_loader backend; + if (!backend.load(impl->backend_cpu, impl->backends, impl->scheduler, impl->buft_policy, true, scheduler_config, + verbosity)) { + error = "failed to initialize StarVLA PI_v3 backend: " + backend.error(); + return nullptr; + } + impl->mode = backend.mode(); + + PIV3GGUFLoader loader(impl->config, impl->weights); + if (!loader.load(path.c_str(), impl->buft_policy.model_buft, impl->loaded, verbosity)) { + error = loader.error(); + return nullptr; + } + if (impl->loaded.ctx_data == nullptr || impl->loaded.model_buffer == nullptr) { + error = "StarVLA PI_v3 policy GGUF has no tensors"; + return nullptr; + } + ggml_backend_buffer_set_usage(impl->loaded.model_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + impl->timestep_table = timestep_projection_table(impl->config); + impl->action_table = action_time_table(impl->config); + if (verbosity >= 1) { + std::fprintf(stderr, "%s: backend=%s qwen=%d width=%d layers=%d horizon=%d action_dim=%d profiles=%zu\n", + __func__, backend_mode_name(impl->mode), impl->config.qwen_hidden_dim, impl->config.dit_width, + impl->config.block_count, impl->config.horizon, impl->config.action_dim, + impl->config.normalization.profiles.size()); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new PIV3Policy(std::move(impl))); +} + +bool PIV3Policy::evaluate(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error) { + return evaluate_internal(qwen_hidden_states, hidden_element_count, qwen_attention_mask, mask_element_count, + initial_noise, noise_element_count, normalized_actions, error); +} + +bool PIV3Policy::evaluate_internal(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, + const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error) { + normalized_actions.clear(); + error.clear(); + if (impl_ == nullptr || impl_->scheduler == nullptr) { + error = "StarVLA PI_v3 policy is not initialized"; + return false; + } + const size_t hidden_width = static_cast(impl_->config.qwen_hidden_dim); + const size_t layer_count = static_cast(impl_->config.qwen_layer_count); + if (qwen_hidden_states == nullptr || qwen_attention_mask == nullptr || initial_noise == nullptr || + mask_element_count == 0 || mask_element_count > static_cast(std::numeric_limits::max()) || + mask_element_count > std::numeric_limits::max() / hidden_width || + mask_element_count * hidden_width > std::numeric_limits::max() / layer_count || + hidden_element_count != mask_element_count * hidden_width * layer_count) { + error = "StarVLA PI_v3 layerwise Qwen conditioning tensor or attention mask has an incompatible shape"; + return false; + } + const size_t expected_noise = static_cast(impl_->config.horizon) * impl_->config.action_dim; + if (noise_element_count != expected_noise) { + error = "StarVLA PI_v3 initial-noise tensor has an incompatible shape"; + return false; + } + if (std::any_of(qwen_hidden_states, qwen_hidden_states + hidden_element_count, + [](float value) { return !std::isfinite(value); }) || + std::any_of(initial_noise, initial_noise + noise_element_count, + [](float value) { return !std::isfinite(value); })) { + error = "StarVLA PI_v3 conditioning and initial noise must be finite"; + return false; + } + bool has_valid_token = false; + for (size_t token = 0; token < mask_element_count; ++token) { + if (qwen_attention_mask[token] > 1) { + error = "StarVLA PI_v3 attention mask values must be zero or one"; + return false; + } + has_valid_token = has_valid_token || qwen_attention_mask[token] != 0; + } + if (!has_valid_token) { + error = "StarVLA PI_v3 attention mask must contain at least one valid token"; + return false; + } + + try { + if (impl_->graph == nullptr || impl_->conditioning_token_count != mask_element_count) { + impl_->build_graph(mask_element_count); + } + } catch (const std::exception & exception) { + error = exception.what(); + return false; + } + + const int query_count = impl_->config.no_state_sequence_length; + const int padded_queries = GGML_PAD(query_count, kKQMaskPad); + std::vector additive_mask(mask_element_count * static_cast(padded_queries), + -std::numeric_limits::infinity()); + for (int query = 0; query < query_count; ++query) { + float * row = additive_mask.data() + static_cast(query) * mask_element_count; + for (size_t token = 0; token < mask_element_count; ++token) { + row[token] = qwen_attention_mask[token] != 0 ? 0.0f : -std::numeric_limits::infinity(); + } + } + + ggml_backend_tensor_set(impl_->hidden_input, qwen_hidden_states, 0, hidden_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->cross_mask_input, additive_mask.data(), 0, additive_mask.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->noise_input, initial_noise, 0, noise_element_count * sizeof(float)); + ggml_backend_tensor_set(impl_->timestep_projection_input, impl_->timestep_table.data(), 0, + impl_->timestep_table.size() * sizeof(float)); + ggml_backend_tensor_set(impl_->action_time_input, impl_->action_table.data(), 0, + impl_->action_table.size() * sizeof(float)); + const float one = 1.0f; + ggml_backend_tensor_set(impl_->scalar_one_input, &one, 0, sizeof(one)); + set_backend_threads(impl_->backends, impl_->n_threads); + if (ggml_backend_sched_graph_compute(impl_->scheduler, impl_->graph) != GGML_STATUS_SUCCESS) { + error = "StarVLA PI_v3 graph compute failed"; + return false; + } + + normalized_actions.resize(expected_noise); + ggml_backend_tensor_get(impl_->output, normalized_actions.data(), 0, expected_noise * sizeof(float)); + if (std::any_of(normalized_actions.begin(), normalized_actions.end(), + [](float value) { return !std::isfinite(value); })) { + normalized_actions.clear(); + error = "StarVLA PI_v3 graph produced non-finite actions"; + return false; + } + return true; +} + +bool PIV3Policy::unnormalize(const std::vector & normalized_actions, const std::string & profile_key_value, + std::vector & actions, std::string & error) const { + if (impl_ == nullptr) { + actions.clear(); + error = "StarVLA PI_v3 policy is not initialized"; + return false; + } + return denormalize_actions(impl_->config.normalization, profile_key_value, normalized_actions, + impl_->config.horizon, impl_->config.action_dim, actions, error); +} + +const PIV3PolicyConfig & PIV3Policy::config() const { + if (impl_ == nullptr) { + throw std::runtime_error("StarVLA PI_v3 policy is not initialized"); + } + return impl_->config; +} + +const char * PIV3Policy::backend_name() const { + return impl_ != nullptr ? backend_mode_name(impl_->mode) : "unknown"; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/pi_v3_policy.h b/src/models/starvla/pi_v3_policy.h new file mode 100644 index 0000000..863e1fd --- /dev/null +++ b/src/models/starvla/pi_v3_policy.h @@ -0,0 +1,92 @@ +#pragma once + +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +struct PIV3PolicyConfig { + std::string backbone_arch; + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + + int qwen_hidden_dim = 0; + int qwen_input_embedding_dim = 0; + int qwen_layer_count = 0; + int qwen_vocab_size = 0; + std::string cot_template; + int image_count = 0; + std::vector image_names; + int image_processor_min_pixels = 0; + int image_processor_max_pixels = 0; + int image_patch_size = 0; + int image_spatial_merge_size = 0; + int image_min_token_count = 0; + int image_max_token_count = 0; + + int dit_width = 0; + int block_count = 0; + int projector_count = 0; + int attention_head_count = 0; + int attention_head_dim = 0; + int feed_forward_dim = 0; + int mlp_hidden_dim = 0; + int action_dim = 0; + int horizon = 0; + int future_token_count = 0; + int action_position_count = 0; + int no_state_sequence_length = 0; + int timestep_projection_dim = 0; + int num_timestep_buckets = 0; + int num_inference_timesteps = 0; + float ada_norm_epsilon = 0.0f; + float projector_norm_epsilon = 0.0f; + float euler_dt = 0.0f; + std::vector qwen_hidden_tuple_indices; + std::vector timestep_ids; + NormalizationConfig normalization; +}; + +class PIV3Policy { + public: + ~PIV3Policy(); + + PIV3Policy(const PIV3Policy &) = delete; + PIV3Policy & operator=(const PIV3Policy &) = delete; + + static std::unique_ptr load(const std::string & path, int n_threads, int verbosity, + std::string & error); + + // qwen_hidden_states is layer-major + // [qwen_layer_count, token_count, qwen_hidden_dim]. Each layer has the + // same full-chat token sequence. Non-zero mask entries participate in + // every cross-attention block. The released checkpoint has no raw-state + // runtime path. initial_noise is token-major [horizon, action_dim]. + bool evaluate(const float * qwen_hidden_states, size_t hidden_element_count, const uint8_t * qwen_attention_mask, + size_t mask_element_count, const float * initial_noise, size_t noise_element_count, + std::vector & normalized_actions, std::string & error); + bool unnormalize(const std::vector & normalized_actions, const std::string & profile_key, + std::vector & actions, std::string & error) const; + + const PIV3PolicyConfig & config() const; + const char * backend_name() const; + + private: + struct Impl; + + explicit PIV3Policy(std::unique_ptr impl); + + bool evaluate_internal(const float * qwen_hidden_states, size_t hidden_element_count, + const uint8_t * qwen_attention_mask, size_t mask_element_count, const float * initial_noise, + size_t noise_element_count, std::vector & normalized_actions, std::string & error); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/policy_gguf.h b/src/models/starvla/policy_gguf.h new file mode 100644 index 0000000..9de19d0 --- /dev/null +++ b/src/models/starvla/policy_gguf.h @@ -0,0 +1,144 @@ +#pragma once + +#include "ggml.h" +#include "gguf.h" +#include "models/starvla/normalization.h" + +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla::detail { + +inline int require_key(gguf_context * gguf, const char * key, gguf_type type) { + const int index = gguf_find_key(gguf, key); + if (index < 0) { + throw std::runtime_error(std::string("missing required StarVLA GGUF metadata: ") + key); + } + if (gguf_get_kv_type(gguf, index) != type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF metadata type: ") + key); + } + return index; +} + +inline std::string require_string(gguf_context * gguf, const char * key) { + return gguf_get_val_str(gguf, require_key(gguf, key, GGUF_TYPE_STRING)); +} + +inline int32_t require_i32(gguf_context * gguf, const char * key) { + return gguf_get_val_i32(gguf, require_key(gguf, key, GGUF_TYPE_INT32)); +} + +inline float require_f32(gguf_context * gguf, const char * key) { + return gguf_get_val_f32(gguf, require_key(gguf, key, GGUF_TYPE_FLOAT32)); +} + +inline bool require_bool(gguf_context * gguf, const char * key) { + return gguf_get_val_bool(gguf, require_key(gguf, key, GGUF_TYPE_BOOL)); +} + +inline int require_array(gguf_context * gguf, const char * key, gguf_type element_type) { + const int index = require_key(gguf, key, GGUF_TYPE_ARRAY); + if (gguf_get_arr_type(gguf, index) != element_type) { + throw std::runtime_error(std::string("invalid StarVLA GGUF array element type: ") + key); + } + return index; +} + +inline std::vector require_string_array(gguf_context * gguf, const char * key) { + const int index = require_array(gguf, key, GGUF_TYPE_STRING); + const size_t count = gguf_get_arr_n(gguf, index); + std::vector values; + values.reserve(count); + for (size_t i = 0; i < count; ++i) { + values.emplace_back(gguf_get_arr_str(gguf, index, i)); + } + return values; +} + +template +inline std::vector require_numeric_array(gguf_context * gguf, const char * key, gguf_type type) { + const int index = require_array(gguf, key, type); + const size_t count = gguf_get_arr_n(gguf, index); + if (count == 0) { + return {}; + } + const auto * data = static_cast(gguf_get_arr_data(gguf, index)); + if (data == nullptr) { + throw std::runtime_error(std::string("missing StarVLA GGUF array data: ") + key); + } + return std::vector(data, data + count); +} + +inline std::vector require_i32_array(gguf_context * gguf, const char * key) { + return require_numeric_array(gguf, key, GGUF_TYPE_INT32); +} + +inline std::vector require_f32_array(gguf_context * gguf, const char * key) { + return require_numeric_array(gguf, key, GGUF_TYPE_FLOAT32); +} + +inline std::vector require_bool_array(gguf_context * gguf, const char * key) { + const auto raw = require_numeric_array(gguf, key, GGUF_TYPE_BOOL); + std::vector values(raw.size()); + for (size_t i = 0; i < raw.size(); ++i) { + values[i] = raw[i] != 0 ? 1 : 0; + } + return values; +} + +inline std::string profile_key(int index, const char * suffix) { + return "starvla.normalization.profile." + std::to_string(index) + "." + suffix; +} + +inline NormalizationConfig require_normalization(gguf_context * gguf, int action_dim) { + NormalizationConfig config; + config.clip_actions = require_bool(gguf, "starvla.normalization.clip_actions"); + config.binary_threshold = require_f32(gguf, "starvla.normalization.binary_threshold"); + config.binary_comparison = require_string(gguf, "starvla.normalization.binary_comparison"); + config.continuous_dimensions = require_i32_array(gguf, "starvla.action.continuous_dimensions"); + config.binary_dimensions = require_i32_array(gguf, "starvla.action.binary_dimensions"); + + const int profile_count = require_i32(gguf, "starvla.normalization.profile_count"); + const auto keys = require_string_array(gguf, "starvla.normalization.profile_keys"); + if (profile_count <= 0 || keys.size() != static_cast(profile_count)) { + throw std::runtime_error("StarVLA normalization profile count is inconsistent"); + } + config.default_profile_key = keys.front(); + config.profiles.reserve(static_cast(profile_count)); + for (int index = 0; index < profile_count; ++index) { + NormalizationProfile profile; + profile.key = require_string(gguf, profile_key(index, "key").c_str()); + profile.action_q01 = require_f32_array(gguf, profile_key(index, "action_q01").c_str()); + profile.action_q99 = require_f32_array(gguf, profile_key(index, "action_q99").c_str()); + profile.action_mask = require_bool_array(gguf, profile_key(index, "action_mask").c_str()); + if (profile.key != keys[static_cast(index)]) { + throw std::runtime_error("StarVLA normalization profile order is inconsistent"); + } + config.profiles.push_back(std::move(profile)); + } + + std::string error; + if (!validate_normalization_config(config, action_dim, error)) { + throw std::runtime_error(error); + } + return config; +} + +inline bool has_shape(const ggml_tensor * tensor, std::initializer_list expected) { + if (tensor == nullptr || static_cast(ggml_n_dims(tensor)) != expected.size()) { + return false; + } + size_t dimension = 0; + for (const int64_t value : expected) { + if (tensor->ne[dimension++] != value) { + return false; + } + } + return true; +} + +} // namespace robotcpp::starvla::detail diff --git a/src/models/starvla/qwen3vl_bridge.cpp b/src/models/starvla/qwen3vl_bridge.cpp new file mode 100644 index 0000000..2c9bd12 --- /dev/null +++ b/src/models/starvla/qwen3vl_bridge.cpp @@ -0,0 +1,1560 @@ +#include "models/starvla/qwen3vl_bridge.h" + +#ifdef ROBOTCPP_STARVLA_CUDA +#include "models/starvla/qwen_bf16_round_cuda.h" +#endif + +#include "ggml.h" +#include "gguf.h" +#include "llama.h" +#include "llama-model.h" +#include "mtmd.h" +#include "models/starvla/oft_prompt.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +bool qwen_vl_resolve_architecture(const std::string & text_architecture, const std::string & projector_type, + QwenVLArchitecture & architecture, std::string & error) { + architecture = QwenVLArchitecture::unknown; + error.clear(); + if (text_architecture == "qwen2vl" && projector_type == "qwen2.5vl_merger") { + architecture = QwenVLArchitecture::qwen2_5_vl; + return true; + } + if (text_architecture == "qwen3vl" && projector_type == "qwen3vl_merger") { + architecture = QwenVLArchitecture::qwen3_vl; + return true; + } + if (text_architecture != "qwen2vl" && text_architecture != "qwen3vl") { + error = "unsupported Qwen-VL text architecture: " + text_architecture; + } else if (projector_type != "qwen2.5vl_merger" && projector_type != "qwen3vl_merger") { + error = "unsupported Qwen-VL projector type: " + projector_type; + } else { + error = "Qwen-VL text and mmproj architectures do not match"; + } + return false; +} + +const char * qwen_vl_architecture_name(QwenVLArchitecture architecture) { + switch (architecture) { + case QwenVLArchitecture::qwen2_5_vl: + return "qwen2.5-vl"; + case QwenVLArchitecture::qwen3_vl: + return "qwen3-vl"; + case QwenVLArchitecture::unknown: + break; + } + return "unknown"; +} + +bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept { + return name != nullptr && (std::strcmp(name, "result_norm") == 0 || std::strcmp(name, "result_embd_pooled") == 0); +} + +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, int decoder_layer_count, int deepstack_layer_count, + int32_t hidden_tuple_index, QwenVLHiddenStateSource & source, std::string & error) { + source = QwenVLHiddenStateSource{}; + error.clear(); + if (decoder_layer_count <= 0 || hidden_tuple_index <= 0 || hidden_tuple_index > decoder_layer_count) { + error = "Qwen-VL hidden-state tuple index is out of range"; + return false; + } + if (architecture == QwenVLArchitecture::qwen2_5_vl) { + if (deepstack_layer_count != 0) { + error = "Qwen2.5-VL hidden-state profile cannot contain DeepStack"; + return false; + } + if (hidden_tuple_index == decoder_layer_count) { + source.kind = QwenVLHiddenStateSourceKind::final_norm; + source.layer = -1; + } else { + source.kind = QwenVLHiddenStateSourceKind::decoder_output; + source.layer = hidden_tuple_index - 1; + } + return true; + } + if (architecture == QwenVLArchitecture::qwen3_vl) { + if (deepstack_layer_count <= 0 || deepstack_layer_count > decoder_layer_count) { + error = "Qwen3-VL DeepStack layer count is incompatible with the model"; + return false; + } + source.kind = hidden_tuple_index <= deepstack_layer_count ? QwenVLHiddenStateSourceKind::deepstack_output + : QwenVLHiddenStateSourceKind::decoder_output; + source.layer = hidden_tuple_index - 1; + return true; + } + error = "Qwen-VL hidden-state architecture is unknown"; + return false; +} + +bool qwen_vl_select_repetition_penalized_top1(const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error) { + token = -1; + error.clear(); + if (logits == nullptr || vocab_size == 0 || vocab_size > static_cast(INT32_MAX) || + !std::isfinite(repetition_penalty) || repetition_penalty <= 0.0f) { + error = "Qwen-VL generation selector received an invalid contract"; + return false; + } + + std::vector repeated(vocab_size, uint8_t{0}); + for (int32_t value : full_sequence) { + if (value < 0 || static_cast(value) >= vocab_size) { + error = "Qwen-VL generated sequence contains an out-of-vocabulary token"; + return false; + } + repeated[static_cast(value)] = 1; + } + + float best = -std::numeric_limits::infinity(); + int32_t best_token = -1; + for (size_t index = 0; index < vocab_size; ++index) { + float score = logits[index]; + if (std::isnan(score)) { + error = "Qwen-VL generation logits contain NaN"; + return false; + } + if (repeated[index] != 0) { + score = score < 0.0f ? score * repetition_penalty : score / repetition_penalty; + } + // torch.argmax returns the first index on ties. + if (best_token < 0 || score > best) { + best = score; + best_token = static_cast(index); + } + } + if (best_token < 0) { + error = "Qwen-VL generation selector did not produce a token"; + return false; + } + token = best_token; + return true; +} + +namespace { + +void quiet_mtmd_log_callback(ggml_log_level level, const char * text, void * user_data) { + (void)user_data; + if (level == GGML_LOG_LEVEL_ERROR) { + std::fputs(text, stderr); + } +} + +struct PreparedMultimodalBatch { + size_t token_count = 0; + llama_pos position_count = 0; + std::vector embeddings; + std::vector positions; + std::vector sequence_counts; + std::vector sequence_values; + std::vector sequences; + std::vector outputs; + std::vector token_ids; + + llama_batch view() { + return { + static_cast(token_count), + nullptr, + embeddings.data(), + positions.data(), + sequence_counts.data(), + sequences.data(), + outputs.data(), + }; + } +}; + +struct BackendPlacement { + bool accelerator_compute = false; + bool cpu_compute = false; +}; + +struct LayerCapture { + BackendPlacement * placement = nullptr; + bool enabled = false; + bool bf16_residual_layer_boundaries = false; + size_t expected_deepstack_layer_count = 0; + size_t token_count = 0; + size_t hidden_size = 0; + std::vector layer_to_slot; + std::vector deepstack_to_slot; + int result_norm_slot = -1; + std::vector values; + std::vector seen; + std::vector rounded_layers; + std::vector rounded_deepstack_layers; +#ifdef ROBOTCPP_STARVLA_CUDA + QwenBF16CaptureCuda cuda_capture; + size_t cuda_captured = 0; +#endif + std::string error; + + void disable() { + enabled = false; + token_count = 0; + hidden_size = 0; + layer_to_slot.clear(); + deepstack_to_slot.clear(); + result_norm_slot = -1; + values.clear(); + seen.clear(); + rounded_layers.clear(); + rounded_deepstack_layers.clear(); +#ifdef ROBOTCPP_STARVLA_CUDA + cuda_captured = 0; +#endif + error.clear(); + } +}; + +bool finish_layer_capture(LayerCapture & capture) { +#ifdef ROBOTCPP_STARVLA_CUDA + if (capture.cuda_captured == 0) { + return true; + } + if (capture.cuda_captured != capture.seen.size()) { + capture.error = "Qwen-VL hidden states span multiple backends"; + return false; + } + return qwen_bf16_capture_download_cuda(capture.cuda_capture, capture.values.data(), capture.values.size(), + capture.error); +#endif + return true; +} + +void begin_layer_boundary_tracking(LayerCapture & capture, size_t layer_count) { + if (!capture.bf16_residual_layer_boundaries) { + capture.rounded_layers.clear(); + capture.rounded_deepstack_layers.clear(); + return; + } + capture.rounded_layers.assign(layer_count, uint8_t{0}); + capture.rounded_deepstack_layers.assign(capture.expected_deepstack_layer_count, uint8_t{0}); +} + +bool validate_layer_boundary_tracking(const LayerCapture & capture, std::string & error) { + if (!capture.bf16_residual_layer_boundaries) { + return true; + } + if (capture.rounded_layers.size() != capture.layer_to_slot.size() || + capture.rounded_deepstack_layers.size() != capture.expected_deepstack_layer_count || + std::any_of(capture.rounded_layers.begin(), capture.rounded_layers.end(), + [](uint8_t seen) { return seen != 1; }) || + std::any_of(capture.rounded_deepstack_layers.begin(), capture.rounded_deepstack_layers.end(), + [](uint8_t seen) { return seen != 1; })) { + error = "Qwen3-VL BF16 residual-boundary roundtrip did not observe " + "every expected l_out/deepstack_out exactly once"; + return false; + } + return true; +} + +bool observe_backend_placement(ggml_tensor * tensor, bool ask, void * user_data) { + if (!ask || tensor == nullptr || tensor->op == GGML_OP_NONE || tensor->buffer == nullptr || user_data == nullptr) { + return false; + } + auto * placement = static_cast(user_data); + ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(tensor->buffer); + ggml_backend_dev_t device = buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); + if (device == nullptr) { + return false; + } + switch (ggml_backend_dev_type(device)) { + case GGML_BACKEND_DEVICE_TYPE_GPU: + case GGML_BACKEND_DEVICE_TYPE_IGPU: + case GGML_BACKEND_DEVICE_TYPE_ACCEL: + placement->accelerator_compute = true; + break; + case GGML_BACKEND_DEVICE_TYPE_CPU: + placement->cpu_compute = true; + break; + case GGML_BACKEND_DEVICE_TYPE_META: + break; + } + return false; +} + +int indexed_output_index(const char * name, const char * prefix) { + const size_t prefix_size = std::strlen(prefix); + if (name == nullptr || std::strncmp(name, prefix, prefix_size) != 0) { + return -1; + } + const char * number = name + prefix_size; + if (*number == '\0') { + return -1; + } + errno = 0; + char * end = nullptr; + const long parsed = std::strtol(number, &end, 10); + if (errno != 0 || end == number || *end != '\0' || parsed < 0 || parsed > INT_MAX) { + return -1; + } + return static_cast(parsed); +} + +bool observe_text_and_capture_layers(ggml_tensor * tensor, bool ask, void * user_data) { + auto * capture = static_cast(user_data); + if (capture == nullptr) { + return false; + } + observe_backend_placement(tensor, ask, capture->placement); + if (!capture->enabled || tensor == nullptr) { + return false; + } + + int slot = -1; + const int deepstack_layer = indexed_output_index(tensor->name, "deepstack_out-"); + const int layer = indexed_output_index(tensor->name, "l_out-"); + const bool is_result_norm = qwen_vl_is_final_norm_tensor_name(tensor->name); + const bool valid_deepstack_layer = + deepstack_layer >= 0 && static_cast(deepstack_layer) < capture->deepstack_to_slot.size(); + const bool valid_layer = layer >= 0 && static_cast(layer) < capture->layer_to_slot.size(); + if (valid_deepstack_layer) { + slot = capture->deepstack_to_slot[static_cast(deepstack_layer)]; + } else if (valid_layer) { + slot = capture->layer_to_slot[static_cast(layer)]; + } else if (is_result_norm) { + slot = capture->result_norm_slot; + } + const bool round_layer = capture->bf16_residual_layer_boundaries && valid_layer; + const bool round_deepstack = capture->bf16_residual_layer_boundaries && valid_deepstack_layer && + static_cast(deepstack_layer) < capture->expected_deepstack_layer_count; + if (slot < 0 && !round_layer && !round_deepstack) { + return false; + } + if (ask) { + return true; + } + + try { + if (slot >= 0) { + if (static_cast(slot) >= capture->seen.size()) { + capture->error = "Qwen3-VL layer-capture slot is out of range"; + return false; + } + if (capture->seen[static_cast(slot)] != 0) { + capture->error = "Qwen3-VL emitted a requested hidden-state output more than once"; + return false; + } + } + if (round_layer && (static_cast(layer) >= capture->rounded_layers.size() || + capture->rounded_layers[static_cast(layer)] != 0)) { + capture->error = "Qwen3-VL l_out BF16 roundtrip index is invalid or repeated"; + return false; + } + if (round_deepstack && (static_cast(deepstack_layer) >= capture->rounded_deepstack_layers.size() || + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] != 0)) { + capture->error = "Qwen3-VL deepstack_out BF16 roundtrip index is invalid or repeated"; + return false; + } + if (!ggml_is_contiguous(tensor) || tensor->ne[0] != static_cast(capture->hidden_size) || + tensor->ne[1] != static_cast(capture->token_count) || tensor->ne[2] != 1 || tensor->ne[3] != 1) { + capture->error = "Qwen3-VL hidden-state output has an incompatible shape or layout"; + return false; + } + if (capture->hidden_size == 0 || + capture->token_count > std::numeric_limits::max() / capture->hidden_size) { + capture->error = "Qwen3-VL hidden-state capture size overflow"; + return false; + } + const size_t count = capture->token_count * capture->hidden_size; + if (count > std::numeric_limits::max() / sizeof(float) || + (slot >= 0 && (count == 0 || static_cast(slot) >= capture->values.size() / count))) { + capture->error = "Qwen3-VL hidden-state capture byte range is invalid"; + return false; + } + bool rounded_on_device = false; + bool captured_on_device = false; +#ifdef ROBOTCPP_STARVLA_CUDA + if ((round_layer || round_deepstack) && tensor->type == GGML_TYPE_F32) { + const QwenBF16RoundStatus status = qwen_bf16_round_cuda(tensor, count, capture->error); + if (status == QwenBF16RoundStatus::error) { + return false; + } + rounded_on_device = status == QwenBF16RoundStatus::success; + } + if (slot >= 0 && capture->seen.size() > 1 && tensor->type == GGML_TYPE_F32) { + const QwenBF16RoundStatus status = + qwen_bf16_capture_cuda(tensor, count, static_cast(slot) * count, capture->values.size(), + capture->cuda_capture, capture->error); + if (status == QwenBF16RoundStatus::error) { + return false; + } + captured_on_device = status == QwenBF16RoundStatus::success; + if (captured_on_device) { + ++capture->cuda_captured; + capture->seen[static_cast(slot)] = 1; + } + } +#endif + if (captured_on_device) { + if (round_layer) { + capture->rounded_layers[static_cast(layer)] = 1; + } else if (round_deepstack) { + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] = 1; + } + return true; + } + std::vector rounded; + if (!rounded_on_device || slot >= 0) { + rounded.resize(count); + } + if (rounded_on_device) { + if (slot >= 0) { + ggml_backend_tensor_get(tensor, rounded.data(), 0, count * sizeof(float)); + } + } else if (tensor->type == GGML_TYPE_F32) { + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(float)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(source[index])); + } + } else if (tensor->type == GGML_TYPE_F16) { + if (round_layer || round_deepstack) { + capture->error = "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + return false; + } + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(ggml_fp16_t)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(ggml_fp16_to_fp32(source[index]))); + } + } else if (tensor->type == GGML_TYPE_BF16) { + if (round_layer || round_deepstack) { + capture->error = "Qwen3-VL BF16 residual-boundary roundtrip requires F32 tensors"; + return false; + } + std::vector source(count); + ggml_backend_tensor_get(tensor, source.data(), 0, count * sizeof(ggml_bf16_t)); + for (size_t index = 0; index < count; ++index) { + rounded[index] = ggml_bf16_to_fp32(source[index]); + } + } else { + capture->error = + std::string("unsupported Qwen3-VL hidden-state output type: ") + ggml_type_name(tensor->type); + return false; + } + if ((round_layer || round_deepstack) && !rounded_on_device) { + ggml_backend_tensor_set(tensor, rounded.data(), 0, count * sizeof(float)); + } + if (round_layer || round_deepstack) { + if (round_layer) { + capture->rounded_layers[static_cast(layer)] = 1; + } else { + capture->rounded_deepstack_layers[static_cast(deepstack_layer)] = 1; + } + } + if (slot >= 0) { + float * destination = capture->values.data() + static_cast(slot) * count; + std::copy(rounded.begin(), rounded.end(), destination); + capture->seen[static_cast(slot)] = 1; + } + return true; + } catch (const std::exception & exception) { + capture->error = std::string("failed to capture Qwen3-VL hidden-state output: ") + exception.what(); + return false; + } catch (...) { + capture->error = "failed to capture Qwen3-VL hidden-state output"; + return false; + } +} + +int32_t decode_and_synchronize(llama_context * context, llama_batch batch) { + const int32_t result = llama_decode(context, batch); + // A cb_eval capture synchronizes only through its selected node. Drain the + // remaining graph tail before disabling capture, clearing KV state, or + // entering a downstream policy graph. + llama_synchronize(context); + return result; +} + +std::string model_metadata(const llama_model * model, const char * key) { + char value[256] = {}; + const int32_t length = llama_model_meta_val_str(model, key, value, sizeof(value)); + if (length < 0 || static_cast(length) >= sizeof(value)) { + throw std::runtime_error(std::string("missing or oversized Qwen GGUF metadata: ") + key); + } + return value; +} + +std::string gguf_string_metadata(const std::string & path, const char * key) { + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = nullptr; + gguf_context * gguf = gguf_init_from_file(path.c_str(), params); + if (gguf == nullptr) { + throw std::runtime_error("failed to read GGUF metadata: " + path); + } + const int64_t index = gguf_find_key(gguf, key); + if (index < 0 || gguf_get_kv_type(gguf, index) != GGUF_TYPE_STRING) { + gguf_free(gguf); + throw std::runtime_error(std::string("missing GGUF string metadata ") + key + ": " + path); + } + const std::string value = gguf_get_val_str(gguf, index); + gguf_free(gguf); + return value; +} + +std::vector tokenize(const llama_vocab * vocab, const std::string & text, bool parse_special) { + const int32_t required = + -llama_tokenize(vocab, text.data(), static_cast(text.size()), nullptr, 0, false, parse_special); + if (required <= 0) { + return {}; + } + std::vector tokens(static_cast(required)); + const int32_t count = llama_tokenize(vocab, text.data(), static_cast(text.size()), tokens.data(), required, + false, parse_special); + if (count != required) { + return {}; + } + return tokens; +} + +std::string apply_chat_template(const llama_model * model, QwenVLArchitecture architecture, + const std::string & content) { + const char * chat_template = llama_model_chat_template(model, nullptr); + if (chat_template == nullptr || chat_template[0] == '\0') { + throw std::runtime_error("Qwen text GGUF has no default chat template"); + } + const llama_chat_message messages[] = { + {"system", "You are a helpful assistant."}, + {"user", content.c_str()}, + }; + const size_t message_offset = architecture == QwenVLArchitecture::qwen2_5_vl ? 0U : 1U; + const size_t message_count = architecture == QwenVLArchitecture::qwen2_5_vl ? 2U : 1U; + const int32_t required = + llama_chat_apply_template(chat_template, messages + message_offset, message_count, true, nullptr, 0); + if (required < 0) { + throw std::runtime_error("failed to size the Qwen chat-template output"); + } + std::vector buffer(static_cast(required) + 1, '\0'); + if (buffer.size() > static_cast(INT32_MAX)) { + throw std::runtime_error("Qwen chat-template output is too large"); + } + const int32_t written = llama_chat_apply_template(chat_template, messages + message_offset, message_count, true, + buffer.data(), static_cast(buffer.size())); + if (written != required) { + throw std::runtime_error("failed to apply the Qwen chat template"); + } + return std::string(buffer.data(), static_cast(written)); +} + +struct PackedImageLayout { + size_t row_bytes = 0; + size_t stride_bytes = 0; + size_t packed_bytes = 0; +}; + +bool validate_image(const Qwen3VLImageView & image, const Qwen3VLBridgeConfig & config, PackedImageLayout & layout, + std::string & error) { + (void)config; + layout = PackedImageLayout{}; + if (image.data == nullptr || image.channels != 3 || image.width <= 0 || image.height <= 0) { + error = "Qwen3-VL bridge requires a non-empty RGB image"; + return false; + } + const size_t width = static_cast(image.width); + const size_t height = static_cast(image.height); + if (width > std::numeric_limits::max() / 3U) { + error = "Qwen3-VL image row size overflow"; + return false; + } + layout.row_bytes = width * 3U; + if (image.stride_bytes < 0 || + (image.stride_bytes > 0 && static_cast(image.stride_bytes) < layout.row_bytes)) { + error = "Qwen3-VL image stride is smaller than a packed RGB row"; + return false; + } + layout.stride_bytes = image.stride_bytes > 0 ? static_cast(image.stride_bytes) : layout.row_bytes; + if (height > std::numeric_limits::max() / layout.row_bytes || + (height > 1U && height - 1U > (std::numeric_limits::max() - layout.row_bytes) / layout.stride_bytes)) { + error = "Qwen3-VL image buffer size overflow"; + return false; + } + layout.packed_bytes = height * layout.row_bytes; + return true; +} + +std::vector pack_image(const Qwen3VLImageView & image, const PackedImageLayout & layout) { + std::vector packed(layout.packed_bytes); + for (int row = 0; row < image.height; ++row) { + std::memcpy(packed.data() + static_cast(row) * layout.row_bytes, + image.data + static_cast(row) * layout.stride_bytes, layout.row_bytes); + } + return packed; +} + +void tokenize_multimodal_prompt(const Qwen3VLBridgeConfig & config, QwenVLArchitecture architecture, + const llama_model * model, mtmd_context * vision, + const std::vector & images, const std::string & instruction, + mtmd::input_chunks & chunks) { + std::vector> packed_images; + packed_images.reserve(images.size()); + mtmd::bitmaps bitmaps; + for (const Qwen3VLImageView & image : images) { + std::string validation_error; + PackedImageLayout layout; + if (!validate_image(image, config, layout, validation_error)) { + throw std::runtime_error(validation_error); + } + packed_images.push_back(pack_image(image, layout)); + bitmaps.entries.emplace_back(static_cast(image.width), static_cast(image.height), + packed_images.back().data()); + if (bitmaps.entries.back().ptr == nullptr) { + throw std::runtime_error("failed to create a Qwen3-VL image bitmap"); + } + } + + const std::string content = build_qwen_media_content(images.size(), instruction, mtmd_default_marker()); + const std::string formatted = apply_chat_template(model, architecture, content); + mtmd_input_text input_text{}; + input_text.text = formatted.c_str(); + input_text.add_special = false; + input_text.parse_special = true; + chunks.ptr.reset(mtmd_input_chunks_init()); + if (chunks.ptr == nullptr) { + throw std::runtime_error("failed to allocate Qwen3-VL multimodal input chunks"); + } + std::vector bitmap_ptrs = bitmaps.c_ptr(); + const int32_t tokenize_result = + mtmd_tokenize(vision, chunks.ptr.get(), &input_text, bitmap_ptrs.data(), bitmap_ptrs.size()); + if (tokenize_result != 0) { + throw std::runtime_error("failed to tokenize the Qwen3-VL multimodal prompt"); + } +} + +const char * compiled_backend_name() { +#if defined(GGML_USE_CUDA) + return "cuda"; +#elif defined(GGML_USE_METAL) + return "metal"; +#else + return "cpu"; +#endif +} + +} // namespace + +struct Qwen3VLBridge::Impl { + Qwen3VLBridgeConfig config; + QwenVLArchitecture architecture = QwenVLArchitecture::unknown; + size_t deepstack_layer_count = 0; + llama_model * model = nullptr; + llama_context * context = nullptr; + mtmd_context * vision = nullptr; + const llama_vocab * vocab = nullptr; + bool backend_initialized = false; + BackendPlacement text_placement; + BackendPlacement vision_placement; + LayerCapture layer_capture; + mutable std::string backend_name = "unknown"; + + void refresh_backend_name() const { + const bool accelerator = text_placement.accelerator_compute && vision_placement.accelerator_compute; + const bool cpu = text_placement.cpu_compute || vision_placement.cpu_compute; + if (accelerator && !cpu) { + backend_name = compiled_backend_name(); + } else if (!text_placement.accelerator_compute && !vision_placement.accelerator_compute && + text_placement.cpu_compute && vision_placement.cpu_compute) { + backend_name = "cpu"; + } else if (text_placement.accelerator_compute || vision_placement.accelerator_compute) { + backend_name = "mixed"; + } else { + backend_name = "unknown"; + } + } + + ~Impl() { + if (vision != nullptr) { + mtmd_free(vision); + vision = nullptr; + } + if (context != nullptr) { + llama_free(context); + context = nullptr; + } + if (model != nullptr) { + llama_model_free(model); + model = nullptr; + } + if (backend_initialized) { + llama_backend_free(); + backend_initialized = false; + } + } +}; + +namespace { + +void copy_token_embedding(const ggml_tensor * token_embeddings, llama_token token, size_t hidden_size, + float * destination) { + if (token_embeddings == nullptr || destination == nullptr || token < 0 || + token_embeddings->ne[0] != static_cast(hidden_size) || token >= token_embeddings->ne[1] || + !ggml_is_contiguous(token_embeddings) || token_embeddings->buffer == nullptr) { + throw std::runtime_error("Qwen3-VL token embedding table is incompatible"); + } + const size_t row_stride = token_embeddings->nb[1]; + if (row_stride == 0 || static_cast(token) > std::numeric_limits::max() / row_stride) { + throw std::runtime_error("Qwen3-VL token embedding row offset overflow"); + } + const size_t row_offset = static_cast(token) * row_stride; + switch (token_embeddings->type) { + case GGML_TYPE_F32: + ggml_backend_tensor_get(token_embeddings, destination, row_offset, hidden_size * sizeof(float)); + break; + case GGML_TYPE_F16: { + std::vector row(hidden_size); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, hidden_size * sizeof(ggml_fp16_t)); + for (size_t index = 0; index < hidden_size; ++index) { + destination[index] = ggml_fp16_to_fp32(row[index]); + } + break; + } + case GGML_TYPE_BF16: { + std::vector row(hidden_size); + ggml_backend_tensor_get(token_embeddings, row.data(), row_offset, hidden_size * sizeof(ggml_bf16_t)); + for (size_t index = 0; index < hidden_size; ++index) { + destination[index] = ggml_bf16_to_fp32(row[index]); + } + break; + } + default: + throw std::runtime_error(std::string("unsupported Qwen3-VL token embedding type: ") + + ggml_type_name(token_embeddings->type)); + } +} + +PreparedMultimodalBatch prepare_multimodal_batch(const Qwen3VLBridgeConfig & config, llama_model * model, + mtmd_context * vision, const mtmd::input_chunks & chunks) { + if (model == nullptr || vision == nullptr || !mtmd_decode_use_mrope(vision)) { + throw std::runtime_error("Qwen3-VL single-batch decode requires M-RoPE components"); + } + + PreparedMultimodalBatch prepared; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + if (chunk_tokens == 0 || chunk_tokens > std::numeric_limits::max() - prepared.token_count) { + throw std::runtime_error("Qwen3-VL multimodal chunk has an invalid token count"); + } + prepared.token_count += chunk_tokens; + } + if (prepared.token_count == 0 || prepared.token_count > static_cast(INT32_MAX)) { + throw std::runtime_error("Qwen3-VL multimodal prompt token count is invalid"); + } + + const size_t hidden_size = static_cast(config.hidden_size); + const size_t input_size = static_cast(config.input_embedding_size); + if (hidden_size == 0 || input_size != static_cast(llama_model_n_embd_inp(model)) || + input_size < hidden_size || prepared.token_count > std::numeric_limits::max() / input_size) { + throw std::runtime_error("Qwen3-VL input embedding dimensions are incompatible"); + } + if (prepared.token_count > std::numeric_limits::max() / 4U) { + throw std::runtime_error("Qwen3-VL M-RoPE position buffer size overflow"); + } + + prepared.embeddings.assign(prepared.token_count * input_size, 0.0f); + prepared.positions.resize(prepared.token_count * 4U); + prepared.sequence_counts.assign(prepared.token_count, 1); + prepared.sequence_values.assign(prepared.token_count, 0); + prepared.sequences.resize(prepared.token_count); + prepared.outputs.assign(prepared.token_count, int8_t{0}); + prepared.token_ids.assign(prepared.token_count, static_cast(-1)); + for (size_t index = 0; index < prepared.token_count; ++index) { + prepared.sequences[index] = &prepared.sequence_values[index]; + } + + const ggml_tensor * token_embeddings = model->get_tensor("token_embd.weight"); + size_t token_offset = 0; + llama_pos position_offset = 0; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + const size_t chunk_tokens = mtmd_input_chunk_get_n_tokens(chunk); + const llama_pos chunk_positions = mtmd_input_chunk_get_n_pos(chunk); + if (chunk_positions <= 0 || position_offset > std::numeric_limits::max() - chunk_positions) { + throw std::runtime_error("Qwen3-VL multimodal positions overflow"); + } + + const mtmd_input_chunk_type type = mtmd_input_chunk_get_type(chunk); + if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) { + size_t text_token_count = 0; + const llama_token * tokens = mtmd_input_chunk_get_tokens_text(chunk, &text_token_count); + if (tokens == nullptr || text_token_count != chunk_tokens || + chunk_positions != static_cast(chunk_tokens)) { + throw std::runtime_error("Qwen3-VL text chunk contract is incompatible"); + } + for (size_t local_index = 0; local_index < chunk_tokens; ++local_index) { + const size_t global_index = token_offset + local_index; + copy_token_embedding(token_embeddings, tokens[local_index], hidden_size, + prepared.embeddings.data() + global_index * input_size); + prepared.token_ids[global_index] = tokens[local_index]; + const llama_pos position = position_offset + static_cast(local_index); + for (size_t axis = 0; axis < 3U; ++axis) { + prepared.positions[axis * prepared.token_count + global_index] = position; + } + prepared.positions[prepared.token_count * 3U + global_index] = 0; + } + } else if (type == MTMD_INPUT_CHUNK_TYPE_IMAGE) { + const mtmd_image_tokens * image_tokens = mtmd_input_chunk_get_tokens_image(chunk); + if (image_tokens == nullptr || mtmd_image_tokens_get_n_tokens(image_tokens) != chunk_tokens) { + throw std::runtime_error("Qwen3-VL image chunk contract is incompatible"); + } + if (mtmd_encode_chunk(vision, chunk) != 0) { + throw std::runtime_error("failed to encode a Qwen3-VL image chunk"); + } + const float * image_embeddings = mtmd_get_output_embd(vision); + if (image_embeddings == nullptr) { + throw std::runtime_error("Qwen3-VL image encoder returned no embeddings"); + } + const size_t image_element_count = chunk_tokens * input_size; + float * destination = prepared.embeddings.data() + token_offset * input_size; + for (size_t element = 0; element < image_element_count; ++element) { + destination[element] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(image_embeddings[element])); + } + for (size_t local_index = 0; local_index < chunk_tokens; ++local_index) { + const size_t global_index = token_offset + local_index; + const mtmd_decoder_pos position = + mtmd_image_tokens_get_decoder_pos(image_tokens, position_offset, local_index); + prepared.positions[global_index] = static_cast(position.t); + prepared.positions[prepared.token_count + global_index] = static_cast(position.y); + prepared.positions[prepared.token_count * 2U + global_index] = static_cast(position.x); + prepared.positions[prepared.token_count * 3U + global_index] = static_cast(position.z); + } + } else { + throw std::runtime_error("Qwen3-VL prompt contains an unsupported media chunk"); + } + token_offset += chunk_tokens; + position_offset += chunk_positions; + } + if (token_offset != prepared.token_count) { + throw std::runtime_error("Qwen3-VL prepared batch token count mismatch"); + } + prepared.position_count = position_offset; + return prepared; +} + +void export_prepared_inputs(const Qwen3VLBridgeConfig & config, const llama_vocab * vocab, + const mtmd::input_chunks & chunks, const PreparedMultimodalBatch & prepared, + std::vector & input_ids, std::vector & attention_mask, + std::vector & image_grid_thw) { + const std::vector image_pad_tokens = tokenize(vocab, "<|image_pad|>", true); + if (image_pad_tokens.size() != 1) { + throw std::runtime_error("Qwen3-VL vocabulary does not expose a unique <|image_pad|> token"); + } + if (config.image_spatial_merge_size <= 0) { + throw std::runtime_error("Qwen3-VL image spatial merge size is invalid"); + } + + input_ids.reserve(prepared.token_ids.size()); + for (llama_token token : prepared.token_ids) { + input_ids.push_back(token < 0 ? static_cast(image_pad_tokens.front()) : static_cast(token)); + } + attention_mask.assign(prepared.token_count, uint8_t{1}); + + image_grid_thw.reserve(static_cast(config.expected_image_count) * 3U); + size_t observed_images = 0; + for (size_t chunk_index = 0; chunk_index < chunks.size(); ++chunk_index) { + const mtmd_input_chunk * chunk = chunks[chunk_index]; + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_IMAGE) { + continue; + } + const mtmd_image_tokens * image_tokens = mtmd_input_chunk_get_tokens_image(chunk); + if (image_tokens == nullptr) { + throw std::runtime_error("Qwen3-VL image chunk has no token grid"); + } + const size_t image_token_count = mtmd_image_tokens_get_n_tokens(image_tokens); + uint32_t max_x = 0; + uint32_t max_y = 0; + for (size_t token = 0; token < image_token_count; ++token) { + const mtmd_decoder_pos position = mtmd_image_tokens_get_decoder_pos(image_tokens, 0, token); + if (position.t != 0 || position.z != 0) { + throw std::runtime_error("Qwen3-VL image token grid does not use the expected M-RoPE layout"); + } + max_x = std::max(max_x, position.x); + max_y = std::max(max_y, position.y); + } + const size_t merged_width = static_cast(max_x) + 1U; + const size_t merged_height = static_cast(max_y) + 1U; + const size_t merge = static_cast(config.image_spatial_merge_size); + if (merged_width == 0 || merged_height == 0 || merged_width > static_cast(INT64_MAX) / merge || + merged_height > static_cast(INT64_MAX) / merge || + merged_width > std::numeric_limits::max() / merged_height || + merged_width * merged_height != image_token_count) { + throw std::runtime_error("Qwen3-VL image token grid is incompatible"); + } + image_grid_thw.push_back(1); + image_grid_thw.push_back(static_cast(merged_height * merge)); + image_grid_thw.push_back(static_cast(merged_width * merge)); + ++observed_images; + } + if (observed_images != static_cast(config.expected_image_count)) { + throw std::runtime_error("Qwen3-VL image grid count does not match the policy"); + } +} + +} // namespace + +Qwen3VLBridge::Qwen3VLBridge(std::unique_ptr impl) : impl_(std::move(impl)) {} + +Qwen3VLBridge::~Qwen3VLBridge() = default; + +std::unique_ptr Qwen3VLBridge::load(const Qwen3VLBridgeConfig & config, std::string & error) { + error.clear(); + const bool action_config_valid = + config.action_token.empty() ? config.action_token_id == -1 : config.action_token_id >= 0; + if (config.text_path.empty() || config.mmproj_path.empty() || config.bundle_uuid.empty() || + config.hidden_size <= 0 || config.input_embedding_size <= 0 || config.vocab_size <= 0 || !action_config_valid || + config.expected_image_count <= 0 || config.image_min_tokens <= 0 || + config.image_max_tokens < config.image_min_tokens || config.image_spatial_merge_size <= 0 || + config.n_ctx <= 0 || config.n_batch <= 0) { + error = "Qwen3-VL bridge configuration is incomplete"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + impl->config = config; + try { + const std::string mmproj_uuid = gguf_string_metadata(config.mmproj_path, "general.source.uuid"); + if (mmproj_uuid != config.bundle_uuid) { + throw std::runtime_error("Qwen3-VL mmproj bundle UUID does not match the policy"); + } + const std::string projector_type = gguf_string_metadata(config.mmproj_path, "clip.projector_type"); + + llama_backend_init(); + impl->backend_initialized = true; + llama_model_params model_params = llama_model_default_params(); + model_params.n_gpu_layers = -1; + impl->model = llama_model_load_from_file(config.text_path.c_str(), model_params); + if (impl->model == nullptr) { + throw std::runtime_error("failed to load Qwen3-VL text GGUF: " + config.text_path); + } + if (model_metadata(impl->model, "general.source.uuid") != config.bundle_uuid) { + throw std::runtime_error("Qwen3-VL text bundle UUID does not match the policy"); + } + std::string profile_error; + if (!qwen_vl_resolve_architecture(model_metadata(impl->model, "general.architecture"), projector_type, + impl->architecture, profile_error)) { + throw std::runtime_error(profile_error); + } + if (llama_model_n_embd_out(impl->model) != config.hidden_size || + llama_model_n_embd_inp(impl->model) != config.input_embedding_size) { + throw std::runtime_error("Qwen3-VL text embedding dimensions do not match the policy"); + } + if (config.input_embedding_size % config.hidden_size != 0) { + throw std::runtime_error("Qwen-VL input embedding width is not an integral hidden-state layout"); + } + const int deepstack_layer_count = config.input_embedding_size / config.hidden_size - 1; + if ((impl->architecture == QwenVLArchitecture::qwen2_5_vl && deepstack_layer_count != 0) || + (impl->architecture == QwenVLArchitecture::qwen3_vl && + (deepstack_layer_count <= 0 || deepstack_layer_count > llama_model_n_layer(impl->model)))) { + throw std::runtime_error("Qwen-VL input embedding layout does not match the detected architecture"); + } + impl->deepstack_layer_count = static_cast(deepstack_layer_count); + for (int layer = 0; layer < llama_model_n_layer(impl->model); ++layer) { + ggml_backend_dev_t device = impl->model->dev_layer(layer); + if (device == nullptr) { + throw std::runtime_error("Qwen3-VL text layer has no assigned backend device"); + } + const enum ggml_backend_dev_type type = ggml_backend_dev_type(device); + if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { + impl->text_placement.cpu_compute = true; + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU || + type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + impl->text_placement.accelerator_compute = true; + } + } + if (ggml_backend_dev_t output_device = impl->model->dev_output()) { + const enum ggml_backend_dev_type type = ggml_backend_dev_type(output_device); + if (type == GGML_BACKEND_DEVICE_TYPE_CPU) { + impl->text_placement.cpu_compute = true; + } else if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU || + type == GGML_BACKEND_DEVICE_TYPE_ACCEL) { + impl->text_placement.accelerator_compute = true; + } + } + + impl->vocab = llama_model_get_vocab(impl->model); + if (impl->vocab == nullptr) { + throw std::runtime_error("Qwen3-VL text GGUF has no vocabulary"); + } + if (llama_vocab_n_tokens(impl->vocab) != config.vocab_size) { + throw std::runtime_error("Qwen3-VL text vocabulary size does not match the policy"); + } + if (!config.action_token.empty()) { + const std::vector action_tokens = tokenize(impl->vocab, config.action_token, true); + if (action_tokens.size() != 1 || action_tokens.front() != config.action_token_id) { + throw std::runtime_error("Qwen3-VL action token mapping does not match the policy metadata"); + } + } + + llama_context_params context_params = llama_context_default_params(); + context_params.n_ctx = static_cast(config.n_ctx); + context_params.n_batch = static_cast(config.n_batch); + // Layer capture expects one complete l_out/deepstack_out tensor per decode call. + context_params.n_ubatch = static_cast(config.n_batch); + context_params.n_threads = config.n_threads; + context_params.n_threads_batch = config.n_threads; + context_params.pooling_type = LLAMA_POOLING_TYPE_NONE; + context_params.embeddings = false; + // Match the official Qwen3-VL BF16 inference cache instead of llama's F16 default. + context_params.type_k = GGML_TYPE_BF16; + context_params.type_v = GGML_TYPE_BF16; + context_params.flash_attn_type = + config.flash_text_attention ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + impl->layer_capture.placement = &impl->text_placement; + impl->layer_capture.bf16_residual_layer_boundaries = config.bf16_residual_layer_boundaries; + if (config.bf16_residual_layer_boundaries) { + impl->layer_capture.expected_deepstack_layer_count = impl->deepstack_layer_count; + } + context_params.cb_eval = observe_text_and_capture_layers; + context_params.cb_eval_user_data = &impl->layer_capture; + impl->context = llama_init_from_model(impl->model, context_params); + if (impl->context == nullptr) { + throw std::runtime_error("failed to create Qwen3-VL text context"); + } + if (config.disable_text_backend_native_graphs) { + llama_set_backend_native_graphs_enabled(impl->context, false); + } + if (llama_n_batch(impl->context) != static_cast(config.n_batch) || + llama_n_ubatch(impl->context) != static_cast(config.n_batch)) { + throw std::runtime_error("Qwen3-VL text context did not preserve the requested batch/ubatch contract"); + } + + mtmd_context_params vision_params = mtmd_context_params_default(); + vision_params.use_gpu = true; + vision_params.print_timings = config.verbosity >= 1; + vision_params.n_threads = config.n_threads; + vision_params.image_min_tokens = config.image_min_tokens; + vision_params.image_max_tokens = config.image_max_tokens; + vision_params.cb_eval = observe_backend_placement; + vision_params.cb_eval_user_data = &impl->vision_placement; + mtmd_log_set(config.verbosity >= 1 ? nullptr : quiet_mtmd_log_callback, nullptr); + impl->vision = mtmd_init_from_file(config.mmproj_path.c_str(), impl->model, vision_params); + if (impl->vision == nullptr) { + throw std::runtime_error("failed to load Qwen3-VL mmproj GGUF: " + config.mmproj_path); + } + if (config.disable_vision_backend_native_graphs) { + mtmd_set_backend_native_graphs_enabled(impl->vision, false); + } + impl->refresh_backend_name(); + + if (config.verbosity >= 1) { + std::fprintf(stderr, + "%s: architecture=%s backend=%s hidden=%d input_embd=%d " + "deepstack=%zu images=%d image_tokens=%d..%d " + "n_ctx=%u n_batch=%u n_ubatch=%u kv=bf16 " + "text_native_graph_disable_requested=%s " + "vision_native_graph_disable_requested=%s\n", + __func__, qwen_vl_architecture_name(impl->architecture), impl->backend_name.c_str(), + llama_model_n_embd_out(impl->model), llama_model_n_embd_inp(impl->model), + impl->deepstack_layer_count, config.expected_image_count, config.image_min_tokens, + config.image_max_tokens, llama_n_ctx(impl->context), llama_n_batch(impl->context), + llama_n_ubatch(impl->context), config.disable_text_backend_native_graphs ? "true" : "false", + config.disable_vision_backend_native_graphs ? "true" : "false"); + } + } catch (const std::exception & exception) { + error = exception.what(); + return nullptr; + } + return std::unique_ptr(new Qwen3VLBridge(std::move(impl))); +} + +bool Qwen3VLBridge::extract_token_embeddings(const std::vector & images, + const std::string & instruction, int32_t token_id, size_t token_count, + std::vector & embeddings, std::string & error) { + embeddings.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (impl_->config.action_token.empty() || impl_->config.action_token_id < 0) { + error = "Qwen3-VL bridge was configured without an action token"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + if (token_count == 0) { + error = "Qwen3-VL requested token embedding count must be positive"; + return false; + } + if (token_id != impl_->config.action_token_id) { + error = "Qwen3-VL requested token ID does not match the policy"; + return false; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + + std::vector matches; + for (size_t index = 0; index < prepared.token_ids.size(); ++index) { + if (prepared.token_ids[index] == token_id) { + matches.push_back(index); + } + } + if (matches.size() < token_count) { + throw std::runtime_error("Qwen3-VL prompt contains fewer target tokens than requested"); + } + matches.erase(matches.begin(), matches.end() - static_cast(token_count)); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + + const int layer_count = llama_model_n_layer(impl_->model); + if (layer_count <= 0) { + throw std::runtime_error("Qwen3-VL model has no decoder layers"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = static_cast(impl_->config.hidden_size); + capture.layer_to_slot.assign(static_cast(layer_count), -1); + capture.deepstack_to_slot.assign(static_cast(layer_count), -1); + capture.result_norm_slot = -1; + if (impl_->architecture == QwenVLArchitecture::qwen2_5_vl) { + capture.result_norm_slot = 0; + } else { + capture.layer_to_slot.back() = 0; + } + capture.values.assign(prepared.token_count * capture.hidden_size, 0.0f); + capture.seen.assign(1, uint8_t{0}); + begin_layer_boundary_tracking(capture, static_cast(layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (capture.seen.size() != 1 || capture.seen.front() == 0) { + throw std::runtime_error("Qwen-VL did not expose the final conditioning output"); + } + + embeddings.resize(token_count * capture.hidden_size); + for (size_t output_index = 0; output_index < matches.size(); ++output_index) { + const float * hidden = capture.values.data() + matches[output_index] * capture.hidden_size; + std::copy_n(hidden, capture.hidden_size, embeddings.data() + output_index * capture.hidden_size); + } + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + embeddings.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::extract_full_hidden_states(const std::vector & images, + const std::string & instruction, std::vector & hidden_states, + std::vector & attention_mask, std::string & error) { + hidden_states.clear(); + attention_mask.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + const size_t hidden_size = static_cast(impl_->config.hidden_size); + if (prepared.token_count > std::numeric_limits::max() / hidden_size) { + throw std::runtime_error("Qwen3-VL hidden-state buffer size overflow"); + } + + const int layer_count = llama_model_n_layer(impl_->model); + if (layer_count <= 0) { + throw std::runtime_error("Qwen3-VL model has no decoder layers"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; + capture.layer_to_slot.assign(static_cast(layer_count), -1); + capture.deepstack_to_slot.assign(static_cast(layer_count), -1); + capture.result_norm_slot = -1; + if (impl_->architecture == QwenVLArchitecture::qwen2_5_vl) { + capture.result_norm_slot = 0; + } else { + capture.layer_to_slot.back() = 0; + } + capture.values.assign(prepared.token_count * hidden_size, 0.0f); + capture.seen.assign(1, uint8_t{0}); + begin_layer_boundary_tracking(capture, static_cast(layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (capture.seen.size() != 1 || capture.seen.front() == 0) { + throw std::runtime_error("Qwen-VL did not expose the final conditioning output"); + } + + hidden_states = std::move(capture.values); + attention_mask.assign(prepared.token_count, uint8_t{1}); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + hidden_states.clear(); + attention_mask.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::extract_layer_hidden_states(const std::vector & images, + const std::string & instruction, + const std::vector & hidden_tuple_indices, + std::vector & hidden_states, + std::vector & attention_mask, std::string & error) { + hidden_states.clear(); + attention_mask.clear(); + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { + error = "Qwen3-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen3-VL image count does not match the policy"; + return false; + } + + const int model_layer_count = llama_model_n_layer(impl_->model); + if (hidden_tuple_indices.empty() || model_layer_count <= 0 || + hidden_tuple_indices.size() > static_cast(model_layer_count)) { + error = "Qwen3-VL requested hidden-state layer set is incompatible with the model"; + return false; + } + std::vector layer_to_slot(static_cast(model_layer_count), -1); + std::vector deepstack_to_slot(static_cast(model_layer_count), -1); + if (impl_->deepstack_layer_count > static_cast(std::numeric_limits::max())) { + error = "Qwen-VL DeepStack layer count exceeds the supported range"; + return false; + } + const int deepstack_layer_count = static_cast(impl_->deepstack_layer_count); + int result_norm_slot = -1; + for (size_t slot = 0; slot < hidden_tuple_indices.size(); ++slot) { + const int32_t tuple_index = hidden_tuple_indices[slot]; + QwenVLHiddenStateSource source; + if (!qwen_vl_hidden_state_source(impl_->architecture, model_layer_count, deepstack_layer_count, tuple_index, + source, error)) { + return false; + } + if (source.kind == QwenVLHiddenStateSourceKind::final_norm) { + if (result_norm_slot >= 0) { + error = "Qwen-VL hidden-state tuple indices must be unique"; + return false; + } + result_norm_slot = static_cast(slot); + continue; + } + if (source.layer < 0 || source.layer >= model_layer_count) { + error = "Qwen-VL hidden-state source layer is out of range"; + return false; + } + std::vector & target = + source.kind == QwenVLHiddenStateSourceKind::deepstack_output ? deepstack_to_slot : layer_to_slot; + if (target[static_cast(source.layer)] >= 0) { + error = "Qwen-VL hidden-state tuple indices must be unique"; + return false; + } + target[static_cast(source.layer)] = static_cast(slot); + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error( + "Qwen3-VL multimodal prompt exceeds n_batch; increase --n-batch for single-batch decode"); + } + if (prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen3-VL multimodal prompt exceeds n_ctx"); + } + + const size_t hidden_size = static_cast(impl_->config.hidden_size); + const size_t requested_layers = hidden_tuple_indices.size(); + if (prepared.token_count > std::numeric_limits::max() / hidden_size || + prepared.token_count * hidden_size > std::numeric_limits::max() / requested_layers) { + throw std::runtime_error("Qwen3-VL layerwise hidden-state buffer size overflow"); + } + LayerCapture & capture = impl_->layer_capture; + capture.enabled = false; + capture.token_count = prepared.token_count; + capture.hidden_size = hidden_size; + capture.layer_to_slot = layer_to_slot; + capture.deepstack_to_slot = deepstack_to_slot; + capture.result_norm_slot = result_norm_slot; + capture.values.assign(requested_layers * prepared.token_count * hidden_size, 0.0f); + capture.seen.assign(requested_layers, uint8_t{0}); + begin_layer_boundary_tracking(capture, static_cast(model_layer_count)); + capture.error.clear(); + capture.enabled = true; + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{1}); + + llama_memory_clear(llama_get_memory(impl_->context), true); + llama_set_embeddings(impl_->context, true); + llama_batch batch = prepared.view(); + const int32_t decode_result = decode_and_synchronize(impl_->context, batch); + capture.enabled = false; + if (decode_result != 0) { + throw std::runtime_error("failed to evaluate the Qwen3-VL multimodal batch"); + } + if (!capture.error.empty()) { + throw std::runtime_error(capture.error); + } + if (!finish_layer_capture(capture)) { + throw std::runtime_error(capture.error); + } + std::string boundary_error; + if (!validate_layer_boundary_tracking(capture, boundary_error)) { + throw std::runtime_error(boundary_error); + } + if (std::any_of(capture.seen.begin(), capture.seen.end(), [](uint8_t seen) { return seen == 0; })) { + throw std::runtime_error("Qwen3-VL did not expose every requested hidden-state output"); + } + + hidden_states = std::move(capture.values); + attention_mask.assign(prepared.token_count, uint8_t{1}); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + hidden_states.clear(); + attention_mask.clear(); + error = exception.what(); + return false; + } +} + +bool Qwen3VLBridge::generate_autoregressive(const std::vector & images, + const std::string & instruction, const QwenVLGenerationConfig & generation, + QwenVLGenerationResult & result, std::string & error) { + result = QwenVLGenerationResult{}; + error.clear(); + if (impl_ == nullptr || impl_->model == nullptr || impl_->context == nullptr || impl_->vision == nullptr || + impl_->vocab == nullptr) { + error = "Qwen-VL bridge is not initialized"; + return false; + } + if (images.size() != static_cast(impl_->config.expected_image_count)) { + error = "Qwen-VL image count does not match the policy"; + return false; + } + if (generation.max_length == 0 || generation.max_length > static_cast(llama_n_ctx(impl_->context)) || + generation.max_length > static_cast(INT32_MAX) || generation.top_k != 1 || + generation.eos_token_ids.empty() || !std::isfinite(generation.repetition_penalty) || + generation.repetition_penalty <= 0.0f) { + error = "Qwen-VL autoregressive generation configuration is incompatible"; + return false; + } + std::vector eos_seen(static_cast(impl_->config.vocab_size), uint8_t{0}); + for (int32_t eos : generation.eos_token_ids) { + if (eos < 0 || eos >= impl_->config.vocab_size || eos_seen[static_cast(eos)] != 0) { + error = "Qwen-VL generation EOS token set is invalid"; + return false; + } + eos_seen[static_cast(eos)] = 1; + } + + try { + mtmd::input_chunks chunks; + tokenize_multimodal_prompt(impl_->config, impl_->architecture, impl_->model, impl_->vision, images, instruction, + chunks); + PreparedMultimodalBatch prepared = prepare_multimodal_batch(impl_->config, impl_->model, impl_->vision, chunks); + if (prepared.token_count > llama_n_batch(impl_->context)) { + throw std::runtime_error("Qwen-VL multimodal prompt exceeds n_batch; increase --n-batch"); + } + if (prepared.token_count > generation.max_length || + prepared.token_count > static_cast(llama_n_ctx(impl_->context)) || + prepared.position_count > static_cast(llama_n_ctx(impl_->context))) { + throw std::runtime_error("Qwen-VL multimodal prompt exceeds the FAST max_length/n_ctx contract"); + } + + std::vector input_ids; + std::vector attention_mask; + std::vector image_grid_thw; + export_prepared_inputs(impl_->config, impl_->vocab, chunks, prepared, input_ids, attention_mask, + image_grid_thw); + if (input_ids.size() != prepared.token_count) { + throw std::runtime_error("Qwen-VL multimodal prompt token export is inconsistent"); + } + result.prompt_token_count = prepared.token_count; + result.full_sequence.reserve(generation.max_length); + for (int64_t input_id : input_ids) { + if (input_id < 0 || input_id >= impl_->config.vocab_size) { + throw std::runtime_error("Qwen-VL multimodal prompt contains an out-of-vocabulary token"); + } + result.full_sequence.push_back(static_cast(input_id)); + } + if (prepared.token_count == generation.max_length) { + impl_->refresh_backend_name(); + return true; + } + + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + std::fill(prepared.outputs.begin(), prepared.outputs.end(), int8_t{0}); + prepared.outputs.back() = 1; + llama_batch prompt_batch = prepared.view(); + if (decode_and_synchronize(impl_->context, prompt_batch) != 0) { + throw std::runtime_error("failed to evaluate the Qwen-VL autoregressive prompt"); + } + + while (result.full_sequence.size() < generation.max_length) { + const float * logits = llama_get_logits_ith(impl_->context, -1); + int32_t next = -1; + std::string selection_error; + if (!qwen_vl_select_repetition_penalized_top1(logits, static_cast(impl_->config.vocab_size), + result.full_sequence, generation.repetition_penalty, next, + selection_error)) { + throw std::runtime_error(selection_error); + } + result.full_sequence.push_back(next); + result.continuation.push_back(next); + if (eos_seen[static_cast(next)] != 0 || result.full_sequence.size() == generation.max_length) { + break; + } + + const size_t generation_index = result.continuation.size() - 1U; + if (generation_index > + static_cast(std::numeric_limits::max() - prepared.position_count)) { + throw std::runtime_error("Qwen-VL autoregressive M-RoPE position overflow"); + } + llama_token token = static_cast(next); + llama_pos position = prepared.position_count + static_cast(generation_index); + int32_t sequence_count = 1; + llama_seq_id sequence_value = 0; + llama_seq_id * sequence = &sequence_value; + int8_t output = 1; + llama_batch token_batch{ + 1, &token, nullptr, &position, &sequence_count, &sequence, &output, + }; + if (decode_and_synchronize(impl_->context, token_batch) != 0) { + throw std::runtime_error("failed to evaluate an incremental Qwen-VL generation token"); + } + } + + llama_memory_clear(llama_get_memory(impl_->context), true); + impl_->refresh_backend_name(); + return true; + } catch (const std::exception & exception) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + result = QwenVLGenerationResult{}; + error = exception.what(); + return false; + } +} + +void Qwen3VLBridge::reset() { + if (impl_ != nullptr && impl_->context != nullptr) { + llama_synchronize(impl_->context); + impl_->layer_capture.disable(); + llama_set_embeddings(impl_->context, false); + llama_memory_clear(llama_get_memory(impl_->context), true); + } +} + +const char * Qwen3VLBridge::backend_name() const { + if (impl_ == nullptr) { + return "unknown"; + } + impl_->refresh_backend_name(); + return impl_->backend_name.c_str(); +} + +QwenVLArchitecture Qwen3VLBridge::architecture() const { + return impl_ != nullptr ? impl_->architecture : QwenVLArchitecture::unknown; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen3vl_bridge.h b/src/models/starvla/qwen3vl_bridge.h new file mode 100644 index 0000000..2c04e5e --- /dev/null +++ b/src/models/starvla/qwen3vl_bridge.h @@ -0,0 +1,160 @@ +#pragma once + +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +enum class QwenVLArchitecture { + unknown, + qwen2_5_vl, + qwen3_vl, +}; + +enum class QwenVLHiddenStateSourceKind { + decoder_output, + deepstack_output, + final_norm, +}; + +struct QwenVLHiddenStateSource { + QwenVLHiddenStateSourceKind kind = QwenVLHiddenStateSourceKind::decoder_output; + int layer = -1; +}; + +// Resolve the paired llama.cpp text and mtmd projector profiles. StarVLA +// supports Qwen2.5-VL and Qwen3-VL only; mismatched text/mmproj files fail +// before either component is evaluated. +bool qwen_vl_resolve_architecture(const std::string & text_architecture, const std::string & projector_type, + QwenVLArchitecture & architecture, std::string & error); + +const char * qwen_vl_architecture_name(QwenVLArchitecture architecture); + +// llama.cpp names the final normalized decoder state `result_norm`, then +// renames the same tensor when embedding output with pooling type NONE is +// enabled. Both names identify the Qwen2.5 hidden_states[-1] boundary. +bool qwen_vl_is_final_norm_tensor_name(const char * name) noexcept; + +// Map one Transformers 4.57 hidden_states tuple index to the corresponding +// llama.cpp graph output. Index zero (the embedding input) is intentionally not +// exposed. Qwen2.5-VL has no DeepStack and its final tuple item is result_norm; +// Qwen3-VL additionally maps the DeepStack outputs exposed by llama.cpp. +bool qwen_vl_hidden_state_source(QwenVLArchitecture architecture, int decoder_layer_count, int deepstack_layer_count, + int32_t hidden_tuple_index, QwenVLHiddenStateSource & source, std::string & error); + +struct Qwen3VLImageView { + const uint8_t * data = nullptr; + int width = 0; + int height = 0; + int channels = 0; + int stride_bytes = 0; +}; + +struct Qwen3VLBridgeConfig { + std::string text_path; + std::string mmproj_path; + std::string bundle_uuid; + int hidden_size = 0; + int input_embedding_size = 0; + int vocab_size = 0; + std::string action_token; + int32_t action_token_id = -1; + int expected_image_count = 0; + int image_min_tokens = 0; + int image_max_tokens = 0; + int image_spatial_merge_size = 0; + int n_ctx = 2048; + int n_batch = 2048; + int n_threads = 0; + int verbosity = 0; + // OFT uses flash attention; other variants require intermediate outputs + // that are only available on the non-flash path. + bool flash_text_attention = false; + // Round each F32 decoder residual output, plus DeepStack outputs when the + // detected architecture has them, through BF16 RNE before it feeds the next + // layer. Intra-layer computation keeps llama.cpp's backend-default profile. + bool bf16_residual_layer_boundaries = false; + // Repeated text decode can rebuild graphs with transient node keys. Disable + // native graph capture/cache for the text context so those keys cannot + // accumulate backend graph instances. Direct graph computation continues. + bool disable_text_backend_native_graphs = false; + // The vision encoder rebuilds its graph for every image. Disable native + // graph capture/cache when its transient graph keys are not stable. + bool disable_vision_backend_native_graphs = false; +}; + +struct QwenVLGenerationConfig { + size_t max_length = 0; + std::vector eos_token_ids; + int top_k = 0; + float repetition_penalty = 0.0f; +}; + +struct QwenVLGenerationResult { + size_t prompt_token_count = 0; + std::vector full_sequence; + std::vector continuation; +}; + +// Implements the deterministic token choice used by the official FAST +// generation profile: Hugging Face repetition penalty over the full sequence, +// followed by top_k=1. Exposed so the generation contract can be tested +// without loading a multi-gigabyte Qwen checkpoint. +bool qwen_vl_select_repetition_penalized_top1(const float * logits, size_t vocab_size, + const std::vector & full_sequence, float repetition_penalty, + int32_t & token, std::string & error); + +class Qwen3VLBridge { + public: + ~Qwen3VLBridge(); + + Qwen3VLBridge(const Qwen3VLBridge &) = delete; + Qwen3VLBridge & operator=(const Qwen3VLBridge &) = delete; + + static std::unique_ptr load(const Qwen3VLBridgeConfig & config, std::string & error); + + bool extract_token_embeddings(const std::vector & images, const std::string & instruction, + int32_t token_id, size_t token_count, std::vector & embeddings, + std::string & error); + + // Full conditioning sequence. Qwen3 uses the outer recorder's raw final + // decoder output (`l_out-(N-1)`); Qwen2.5 uses `result_norm`, matching its + // Transformers hidden_states[-1]. Values are widened from BF16. + bool extract_full_hidden_states(const std::vector & images, const std::string & instruction, + std::vector & hidden_states, std::vector & attention_mask, + std::string & error); + + // hidden_tuple_indices use the pinned Transformers 4.57 convention. Index + // zero is the embedding input and is not exposed. For Qwen3, in-place + // DeepStack aliases make the first D entries `deepstack_out`; remaining + // entries, including N, are raw `l_out`. For Qwen2.5, indices 1..N-1 are + // raw `l_out` and index N is `result_norm`. The result is layer-major + // [requested states, tokens, hidden size]. + bool extract_layer_hidden_states(const std::vector & images, const std::string & instruction, + const std::vector & hidden_tuple_indices, + std::vector & hidden_states, std::vector & attention_mask, + std::string & error); + + // Runs a full multimodal prefill followed by incremental KV-cached text + // decoding. The returned sequence includes the prompt, matching + // Transformers generate(return_dict_in_generate=false). + bool generate_autoregressive(const std::vector & images, const std::string & instruction, + const QwenVLGenerationConfig & generation, QwenVLGenerationResult & result, + std::string & error); + + void reset(); + const char * backend_name() const; + QwenVLArchitecture architecture() const; + + private: + struct Impl; + + explicit Qwen3VLBridge(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen_bf16_round_cuda.cu b/src/models/starvla/qwen_bf16_round_cuda.cu new file mode 100644 index 0000000..2540020 --- /dev/null +++ b/src/models/starvla/qwen_bf16_round_cuda.cu @@ -0,0 +1,218 @@ +#include "models/starvla/qwen_bf16_round_cuda.h" + +#include "ggml-backend.h" +#include "ggml.h" + +#include + +#include + +namespace robotcpp::starvla { +namespace { + +__device__ uint32_t bf16_bits(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x7fffffffU) > 0x7f800000U + ? (bits >> 16) | 64U + : (bits + 0x7fffU + ((bits >> 16) & 1U)) >> 16; +} + +__global__ void round_bf16(float * values, size_t count) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + values[index] = __uint_as_float(bf16_bits(values[index]) << 16); + } +} + +__global__ void capture_bf16(const float * source, float * destination, + size_t count) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index < count) { + destination[index] = __uint_as_float(bf16_bits(source[index]) << 16); + } +} + +const char * cuda_error(cudaError_t status) { + return cudaGetErrorString(status); +} + +cudaError_t select_device(int device, int & previous_device) { + const cudaError_t status = cudaGetDevice(&previous_device); + return status == cudaSuccess ? cudaSetDevice(device) : status; +} + +QwenBF16RoundStatus tensor_device(ggml_tensor * tensor, size_t count, + int & device_id, std::string & error) { + if (tensor == nullptr || tensor->buffer == nullptr || tensor->data == nullptr || + tensor->type != GGML_TYPE_F32 || count == 0) { + error = "Qwen-VL CUDA BF16 operation received an invalid tensor"; + return QwenBF16RoundStatus::error; + } + ggml_backend_buffer_type_t buffer_type = + ggml_backend_buffer_get_type(tensor->buffer); + ggml_backend_dev_t device = + buffer_type == nullptr ? nullptr : ggml_backend_buft_get_device(buffer_type); + if (device == nullptr || + ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_GPU) { + return QwenBF16RoundStatus::unavailable; + } + + cudaPointerAttributes attributes{}; + const cudaError_t status = cudaPointerGetAttributes(&attributes, tensor->data); + if (status == cudaErrorInvalidValue) { + cudaGetLastError(); + return QwenBF16RoundStatus::unavailable; + } + if (status != cudaSuccess) { + error = std::string("failed to inspect Qwen-VL CUDA tensor: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (attributes.type != cudaMemoryTypeDevice && + attributes.type != cudaMemoryTypeManaged) { + return QwenBF16RoundStatus::unavailable; + } + device_id = attributes.device; + return QwenBF16RoundStatus::success; +} + +} // namespace + +QwenBF16CaptureCuda::~QwenBF16CaptureCuda() { + if (data == nullptr) { + return; + } + int previous_device = 0; + if (select_device(device, previous_device) == cudaSuccess) { + cudaFree(data); + cudaSetDevice(previous_device); + } +} + +QwenBF16RoundStatus qwen_bf16_round_cuda(ggml_tensor * tensor, size_t count, + std::string & error) { + error.clear(); + int device = -1; + const QwenBF16RoundStatus available = + tensor_device(tensor, count, device, error); + if (available != QwenBF16RoundStatus::success) { + return available; + } + + int previous_device = 0; + cudaError_t status = select_device(device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + + constexpr int block_size = 256; + const size_t block_count = (count + block_size - 1) / block_size; + round_bf16<<>>( + static_cast(tensor->data), count); + status = cudaGetLastError(); + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess) { + error = std::string("failed to round Qwen-VL residuals on CUDA: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (restore_status != cudaSuccess) { + error = std::string("failed to restore the active CUDA device: ") + + cuda_error(restore_status); + return QwenBF16RoundStatus::error; + } + return QwenBF16RoundStatus::success; +} + +QwenBF16RoundStatus qwen_bf16_capture_cuda( + ggml_tensor * tensor, size_t count, size_t offset, size_t total_count, + QwenBF16CaptureCuda & capture, std::string & error) { + error.clear(); + if (offset > total_count || count > total_count - offset || + total_count > SIZE_MAX / sizeof(float)) { + error = "Qwen-VL CUDA capture range is invalid"; + return QwenBF16RoundStatus::error; + } + int device = -1; + const QwenBF16RoundStatus available = + tensor_device(tensor, count, device, error); + if (available != QwenBF16RoundStatus::success) { + return available; + } + + int previous_device = 0; + cudaError_t status = select_device(device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return QwenBF16RoundStatus::error; + } + if (capture.data != nullptr && + (capture.device != device || capture.capacity < total_count)) { + status = cudaFree(capture.data); + if (status == cudaSuccess) { + capture.data = nullptr; + capture.capacity = 0; + capture.device = -1; + } + } + if (status == cudaSuccess && capture.data == nullptr) { + status = cudaMalloc(&capture.data, total_count * sizeof(float)); + if (status == cudaSuccess) { + capture.capacity = total_count; + capture.device = device; + } + } + constexpr int block_size = 256; + const size_t block_count = (count + block_size - 1) / block_size; + if (status == cudaSuccess) { + capture_bf16<<>>( + static_cast(tensor->data), + static_cast(capture.data) + offset, count); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess || restore_status != cudaSuccess) { + error = std::string("failed to capture Qwen-VL hidden states on CUDA: ") + + cuda_error(status != cudaSuccess ? status : restore_status); + return QwenBF16RoundStatus::error; + } + return QwenBF16RoundStatus::success; +} + +bool qwen_bf16_capture_download_cuda(QwenBF16CaptureCuda & capture, + float * values, size_t count, + std::string & error) { + error.clear(); + if (capture.data == nullptr || values == nullptr || count == 0 || + count > capture.capacity || count > SIZE_MAX / sizeof(float)) { + error = "Qwen-VL CUDA capture download is invalid"; + return false; + } + int previous_device = 0; + cudaError_t status = select_device(capture.device, previous_device); + if (status != cudaSuccess) { + error = std::string("failed to select the Qwen-VL CUDA device: ") + + cuda_error(status); + return false; + } + status = cudaMemcpy(values, capture.data, count * sizeof(float), + cudaMemcpyDeviceToHost); + const cudaError_t restore_status = cudaSetDevice(previous_device); + if (status != cudaSuccess || restore_status != cudaSuccess) { + error = std::string("failed to download Qwen-VL CUDA hidden states: ") + + cuda_error(status != cudaSuccess ? status : restore_status); + return false; + } + return true; +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/qwen_bf16_round_cuda.h b/src/models/starvla/qwen_bf16_round_cuda.h new file mode 100644 index 0000000..985ed97 --- /dev/null +++ b/src/models/starvla/qwen_bf16_round_cuda.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include + +struct ggml_tensor; + +namespace robotcpp::starvla { + +enum class QwenBF16RoundStatus { unavailable, success, error }; + +struct QwenBF16CaptureCuda { + QwenBF16CaptureCuda() = default; + ~QwenBF16CaptureCuda(); + QwenBF16CaptureCuda(const QwenBF16CaptureCuda &) = delete; + QwenBF16CaptureCuda & operator=(const QwenBF16CaptureCuda &) = delete; + + void * data = nullptr; + size_t capacity = 0; + int device = -1; +}; + +QwenBF16RoundStatus qwen_bf16_round_cuda(ggml_tensor * tensor, size_t count, std::string & error); +QwenBF16RoundStatus qwen_bf16_capture_cuda(ggml_tensor * tensor, size_t count, size_t offset, size_t total_count, + QwenBF16CaptureCuda & capture, std::string & error); +bool qwen_bf16_capture_download_cuda(QwenBF16CaptureCuda & capture, float * values, size_t count, std::string & error); + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_engine.cpp b/src/models/starvla/starvla_engine.cpp new file mode 100644 index 0000000..db23ad5 --- /dev/null +++ b/src/models/starvla/starvla_engine.cpp @@ -0,0 +1,987 @@ +#include "models/starvla/starvla_engine.h" + +#include "ggml.h" +#include "gguf.h" +#include "models/starvla/fast_policy.h" +#include "models/starvla/groot_policy.h" +#include "models/starvla/groot_prompt.h" +#include "models/starvla/normalization.h" +#include "models/starvla/oft_image_preprocess.h" +#include "models/starvla/oft_policy.h" +#include "models/starvla/oft_prompt.h" +#include "models/starvla/pi_policy.h" +#include "models/starvla/pi_v3_policy.h" +#include "models/starvla/qwen3vl_bridge.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace robotcpp::starvla { + +const char * starvla_variant_name(StarVLAVariant variant) noexcept { + switch (variant) { + case StarVLAVariant::qwen3_oft: + return "qwen3_oft"; + case StarVLAVariant::qwen3_groot: + return "qwen3_groot"; + case StarVLAVariant::qwen3_pi_v3: + return "qwen3_pi_v3"; + case StarVLAVariant::qwen25_oft: + return "qwen25_oft"; + case StarVLAVariant::qwen25_groot: + return "qwen25_groot"; + case StarVLAVariant::qwen25_pi: + return "qwen25_pi"; + case StarVLAVariant::qwen25_fast: + return "qwen25_fast"; + } + return "unknown"; +} + +const char * starvla_variant_framework(StarVLAVariant variant) noexcept { + switch (variant) { + case StarVLAVariant::qwen3_oft: + case StarVLAVariant::qwen25_oft: + return "oft"; + case StarVLAVariant::qwen3_groot: + case StarVLAVariant::qwen25_groot: + return "groot"; + case StarVLAVariant::qwen3_pi_v3: + return "pi_v3"; + case StarVLAVariant::qwen25_pi: + return "pi"; + case StarVLAVariant::qwen25_fast: + return "fast"; + } + return "unknown"; +} + +bool starvla_variant_from_metadata(const std::string & framework, const std::string & backbone, + StarVLAVariant & variant) noexcept { + if (backbone == "qwen3_vl") { + if (framework == "oft") + variant = StarVLAVariant::qwen3_oft; + else if (framework == "groot") + variant = StarVLAVariant::qwen3_groot; + else if (framework == "pi_v3") + variant = StarVLAVariant::qwen3_pi_v3; + else + return false; + return true; + } + if (backbone == "qwen2_5_vl") { + if (framework == "oft") + variant = StarVLAVariant::qwen25_oft; + else if (framework == "groot") + variant = StarVLAVariant::qwen25_groot; + else if (framework == "pi") + variant = StarVLAVariant::qwen25_pi; + else if (framework == "fast") + variant = StarVLAVariant::qwen25_fast; + else + return false; + return true; + } + return false; +} + +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int kDefaultThreadCount = 4; + +double elapsed_ms(Clock::time_point start, Clock::time_point end) { + return std::chrono::duration(end - start).count(); +} + +bool read_policy_variant(const std::filesystem::path & path, StarVLAVariant & variant, std::string & error) { + gguf_init_params params{}; + params.no_alloc = true; + gguf_context * gguf = gguf_init_from_file(path.string().c_str(), params); + if (gguf == nullptr) { + error = "failed to read StarVLA policy GGUF metadata"; + return false; + } + const auto read_string = [&](const char * key, std::string & value) { + const int index = gguf_find_key(gguf, key); + if (index < 0 || gguf_get_kv_type(gguf, index) != GGUF_TYPE_STRING) { + error = std::string("missing StarVLA policy metadata: ") + key; + return false; + } + value = gguf_get_val_str(gguf, index); + return true; + }; + + std::string framework; + std::string backbone; + const bool valid = read_string("starvla.framework", framework) && read_string("starvla.backbone.arch", backbone); + gguf_free(gguf); + if (!valid) { + return false; + } + + if (starvla_variant_from_metadata(framework, backbone, variant)) { + return true; + } + error = "unsupported StarVLA variant: " + backbone + "/" + framework; + return false; +} + +bool is_plain_basename(const std::string & value) { + if (value.empty() || value.find('\0') != std::string::npos) { + return false; + } + const std::filesystem::path path(value); + return !path.has_root_path() && !path.has_parent_path() && path.filename() == path && value != "." && value != ".."; +} + +bool require_regular_file(const std::filesystem::path & path, const char * label, std::string & error) { + std::error_code status_error; + const bool regular = std::filesystem::is_regular_file(path, status_error); + if (!regular) { + error = std::string("StarVLA ") + label + " is not a regular file: " + path.string(); + if (status_error) { + error += " (" + status_error.message() + ")"; + } + return false; + } + return true; +} + +bool resolve_component_path(const std::string & metadata_filename, const std::string & component_path, + const char * label, std::filesystem::path & resolved, std::string & error) { + if (!is_plain_basename(metadata_filename)) { + error = std::string("StarVLA policy ") + label + " filename must be a plain basename: " + metadata_filename; + return false; + } + + if (component_path.empty() || component_path.find('\0') != std::string::npos) { + error = std::string("StarVLA ") + label + " path is required and must not contain an embedded NUL"; + return false; + } + resolved = std::filesystem::path(component_path); + if (resolved.filename().string() != metadata_filename) { + error = std::string("StarVLA ") + label + " basename must match policy metadata '" + metadata_filename + + "': " + resolved.string(); + return false; + } + return require_regular_file(resolved, label, error); +} + +bool same_file(const std::filesystem::path & lhs, const std::filesystem::path & rhs) { + std::error_code equivalent_error; + return std::filesystem::equivalent(lhs, rhs, equivalent_error) && !equivalent_error; +} + +bool validate_observation(const observation & obs, int image_count, const std::vector & image_names, + bool state_supported, const char * framework, std::string & error) { + const std::string label = std::string("StarVLA ") + framework; + if (obs.images.size() != static_cast(image_count)) { + error = label + " requires exactly " + std::to_string(image_count) + " image(s) in policy order"; + return false; + } + if (obs.task.empty()) { + error = label + " task must not be empty"; + return false; + } + if (obs.task.find('\0') != std::string::npos) { + error = label + " task contains an embedded NUL"; + return false; + } + if (!state_supported && !obs.state.empty()) { + error = label + " released checkpoint does not support state input"; + return false; + } + for (float value : obs.state) { + if (!std::isfinite(value)) { + error = label + " state must contain only finite values"; + return false; + } + } + + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + if (image.name != image_names[i]) { + error = label + " image " + std::to_string(i) + " must be named '" + image_names[i] + "'"; + return false; + } + if (image.data == nullptr || image.width <= 0 || image.height <= 0) { + error = label + " image '" + image.name + "' has invalid data or dimensions"; + return false; + } + if (image.channels != 3) { + error = label + " image '" + image.name + "' must use interleaved RGB channels"; + return false; + } + if (image.width > std::numeric_limits::max() / image.channels) { + error = label + " image '" + image.name + "' row size overflows"; + return false; + } + const int packed_stride = image.width * image.channels; + if (image.stride_bytes < 0 || (image.stride_bytes != 0 && image.stride_bytes < packed_stride)) { + error = label + " image '" + image.name + "' stride is smaller than a packed RGB row"; + return false; + } + } + return true; +} + +template +bool prepare_qwen_images(const observation & obs, const PolicyConfig & config, const char * framework, + std::vector> & processed_images, + std::vector & qwen_images, std::string & error) { + processed_images.clear(); + qwen_images.clear(); + processed_images.resize(obs.images.size()); + qwen_images.reserve(obs.images.size()); + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + int target_width = 0; + int target_height = 0; + int image_token_count = 0; + std::string preprocess_error; + if (!preprocess_qwen3vl_rgb( + image.data, image.width, image.height, image.channels, image.stride_bytes, config.image_patch_size, + config.image_spatial_merge_size, config.image_processor_min_pixels, config.image_processor_max_pixels, + processed_images[i], target_width, target_height, image_token_count, preprocess_error)) { + error = std::string("failed to preprocess StarVLA ") + framework + " image '" + image.name + + "': " + preprocess_error; + return false; + } + const uint64_t expected_bytes = static_cast(target_width) * static_cast(target_height) * 3; + if (expected_bytes != processed_images[i].size() || image_token_count < config.image_min_token_count || + image_token_count > config.image_max_token_count) { + error = std::string("StarVLA ") + framework + " image preprocessor returned an incompatible dynamic grid"; + return false; + } + Qwen3VLImageView view; + view.data = processed_images[i].data(); + view.width = target_width; + view.height = target_height; + view.channels = 3; + view.stride_bytes = target_width * 3; + qwen_images.push_back(view); + } + return true; +} + +bool prepare_pi_qwen_images(const observation & obs, const PIPolicyConfig & config, + std::vector> & pre_resized_images, + std::vector> & processed_images, + std::vector & qwen_images, std::string & error) { + pre_resized_images.clear(); + processed_images.clear(); + qwen_images.clear(); + pre_resized_images.resize(obs.images.size()); + processed_images.resize(obs.images.size()); + qwen_images.reserve(obs.images.size()); + for (size_t i = 0; i < obs.images.size(); ++i) { + const model_image & image = obs.images[i]; + std::string preprocess_error; + if (!resize_torchvision_bicubic_aa_rgb(image.data, image.width, image.height, image.stride_bytes, + config.image_framework_inference_pre_resize_width, + config.image_framework_inference_pre_resize_height, + pre_resized_images[i], preprocess_error)) { + error = "failed to pre-resize StarVLA PI image '" + image.name + "': " + preprocess_error; + return false; + } + + int target_width = 0; + int target_height = 0; + int image_token_count = 0; + if (!preprocess_qwen3vl_rgb(pre_resized_images[i].data(), config.image_framework_inference_pre_resize_width, + config.image_framework_inference_pre_resize_height, 3, + config.image_framework_inference_pre_resize_width * 3, config.image_patch_size, + config.image_spatial_merge_size, config.image_processor_min_pixels, + config.image_processor_max_pixels, processed_images[i], target_width, target_height, + image_token_count, preprocess_error)) { + error = "failed to preprocess StarVLA PI image '" + image.name + "': " + preprocess_error; + return false; + } + const uint64_t expected_bytes = static_cast(target_width) * static_cast(target_height) * 3; + if (expected_bytes != processed_images[i].size() || image_token_count < config.image_min_token_count || + image_token_count > config.image_max_token_count) { + error = "StarVLA PI image preprocessor returned an incompatible dynamic grid"; + return false; + } + Qwen3VLImageView view; + view.data = processed_images[i].data(); + view.width = target_width; + view.height = target_height; + view.channels = 3; + view.stride_bytes = target_width * 3; + qwen_images.push_back(view); + } + return true; +} + +} // namespace + +struct StarVLAEngine::Impl { + StarVLAVariant variant = StarVLAVariant::qwen3_oft; + std::filesystem::path policy_path; + std::filesystem::path text_path; + std::filesystem::path mmproj_path; + std::string normalization_profile_key; + const NormalizationConfig * normalization = nullptr; + std::mt19937_64 noise_rng; + // Destroy the policy scheduler/backends before Qwen releases llama's global backend state. + std::unique_ptr qwen; + std::unique_ptr oft_policy; + std::unique_ptr groot_policy; + std::unique_ptr pi_policy; + std::unique_ptr pi_v3_policy; + std::unique_ptr fast_policy; +}; + +StarVLAEngine::StarVLAEngine(std::unique_ptr impl) : impl_(std::move(impl)) {} + +StarVLAEngine::~StarVLAEngine() = default; + +std::unique_ptr StarVLAEngine::load(const StarVLAEngineConfig & config, std::string & error) { + error.clear(); + if (config.policy_path.empty() || config.policy_path.find('\0') != std::string::npos) { + error = "StarVLA policy path is required and must not contain an embedded NUL"; + return nullptr; + } + if (config.n_ctx <= 0 || config.n_batch <= 0 || config.n_threads < 0) { + error = "StarVLA n_ctx/n_batch must be positive and n_threads must be non-negative"; + return nullptr; + } + + std::unique_ptr impl(new Impl()); + const int effective_threads = config.n_threads > 0 ? config.n_threads : kDefaultThreadCount; + impl->policy_path = std::filesystem::path(config.policy_path); + if (!require_regular_file(impl->policy_path, "policy GGUF", error)) { + return nullptr; + } + if (!read_policy_variant(impl->policy_path, impl->variant, error)) { + return nullptr; + } + const bool is_oft = impl->variant == StarVLAVariant::qwen3_oft || impl->variant == StarVLAVariant::qwen25_oft; + const bool is_groot = impl->variant == StarVLAVariant::qwen3_groot || impl->variant == StarVLAVariant::qwen25_groot; + const bool is_pi = impl->variant == StarVLAVariant::qwen25_pi; + const bool is_pi_v3 = impl->variant == StarVLAVariant::qwen3_pi_v3; + const bool is_fast = impl->variant == StarVLAVariant::qwen25_fast; + const char * framework = starvla_variant_framework(impl->variant); + + std::string bundle_uuid; + std::string text_filename; + std::string mmproj_filename; + std::string qwen_backbone_arch; + int qwen_hidden_size = 0; + int qwen_input_embedding_size = 0; + int qwen_vocab_size = 0; + int image_count = 0; + int image_min_tokens = 0; + int image_max_tokens = 0; + int image_spatial_merge_size = 0; + const NormalizationConfig * normalization = nullptr; + if (is_oft) { + impl->oft_policy = OFTPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->oft_policy == nullptr) { + error = "failed to load StarVLA OFT policy: " + error; + return nullptr; + } + const OFTPolicyConfig & policy = impl->oft_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.input_dim; + qwen_input_embedding_size = policy.input_embedding_dim; + qwen_vocab_size = policy.vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_groot) { + impl->groot_policy = GR00TPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->groot_policy == nullptr) { + error = "failed to load StarVLA GR00T policy: " + error; + return nullptr; + } + const GR00TPolicyConfig & policy = impl->groot_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_pi) { + impl->pi_policy = PIPolicy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->pi_policy == nullptr) { + error = "failed to load StarVLA PI policy: " + error; + return nullptr; + } + const PIPolicyConfig & policy = impl->pi_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else if (is_pi_v3) { + impl->pi_v3_policy = PIV3Policy::load(impl->policy_path.string(), effective_threads, config.verbosity, error); + if (impl->pi_v3_policy == nullptr) { + error = "failed to load StarVLA PI_v3 policy: " + error; + return nullptr; + } + const PIV3PolicyConfig & policy = impl->pi_v3_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + } else { + impl->fast_policy = FastPolicy::load(impl->policy_path.string(), config.verbosity, error); + if (impl->fast_policy == nullptr) { + error = "failed to load StarVLA FAST policy: " + error; + return nullptr; + } + const FastPolicyConfig & policy = impl->fast_policy->config(); + bundle_uuid = policy.bundle_uuid; + text_filename = policy.text_filename; + mmproj_filename = policy.mmproj_filename; + qwen_backbone_arch = policy.backbone_arch; + qwen_hidden_size = policy.qwen_hidden_dim; + qwen_input_embedding_size = policy.qwen_input_embedding_dim; + qwen_vocab_size = policy.qwen_vocab_size; + image_count = policy.image_count; + image_min_tokens = policy.image_min_token_count; + image_max_tokens = policy.image_max_token_count; + image_spatial_merge_size = policy.image_spatial_merge_size; + normalization = &policy.normalization; + if (config.n_ctx < static_cast(policy.generation_max_length)) { + error = "StarVLA FAST --n-ctx must be at least max_length=" + std::to_string(policy.generation_max_length); + return nullptr; + } + } + if (!resolve_component_path(text_filename, config.text_path, "text GGUF", impl->text_path, error) || + !resolve_component_path(mmproj_filename, config.mmproj_path, "mmproj GGUF", impl->mmproj_path, error)) { + return nullptr; + } + if (same_file(impl->policy_path, impl->text_path) || same_file(impl->policy_path, impl->mmproj_path) || + same_file(impl->text_path, impl->mmproj_path)) { + error = "StarVLA policy, text, and mmproj GGUF paths must identify three distinct files"; + return nullptr; + } + + impl->normalization = normalization; + std::string profile_error; + const NormalizationProfile * profile = resolve_normalization_profile(*normalization, "", profile_error); + if (profile == nullptr) { + error = std::string("failed to select StarVLA ") + framework + " normalization profile: " + profile_error; + return nullptr; + } + impl->normalization_profile_key = profile->key; + + Qwen3VLBridgeConfig qwen_config; + qwen_config.text_path = impl->text_path.string(); + qwen_config.mmproj_path = impl->mmproj_path.string(); + qwen_config.bundle_uuid = bundle_uuid; + qwen_config.hidden_size = qwen_hidden_size; + qwen_config.input_embedding_size = qwen_input_embedding_size; + qwen_config.vocab_size = qwen_vocab_size; + if (is_oft) { + qwen_config.action_token = impl->oft_policy->config().prompt.action_token; + qwen_config.action_token_id = impl->oft_policy->config().action_token_id; + } else { + qwen_config.action_token.clear(); + qwen_config.action_token_id = -1; + } + qwen_config.expected_image_count = image_count; + qwen_config.image_min_tokens = image_min_tokens; + qwen_config.image_max_tokens = image_max_tokens; + qwen_config.image_spatial_merge_size = image_spatial_merge_size; + qwen_config.n_ctx = config.n_ctx; + qwen_config.n_batch = config.n_batch; + qwen_config.n_threads = effective_threads; + qwen_config.verbosity = config.verbosity; + qwen_config.flash_text_attention = is_oft || is_pi; + qwen_config.bf16_residual_layer_boundaries = is_groot; + qwen_config.disable_text_backend_native_graphs = true; + qwen_config.disable_vision_backend_native_graphs = true; + if (qwen_config.input_embedding_size <= 0) { + error = "StarVLA Qwen-VL input embedding size is invalid"; + return nullptr; + } + impl->qwen = Qwen3VLBridge::load(qwen_config, error); + if (impl->qwen == nullptr) { + error = "failed to load StarVLA Qwen-VL components: " + error; + return nullptr; + } + const QwenVLArchitecture expected_architecture = + qwen_backbone_arch == "qwen2_5_vl" ? QwenVLArchitecture::qwen2_5_vl : QwenVLArchitecture::qwen3_vl; + if (impl->qwen->architecture() != expected_architecture) { + error = "StarVLA policy backbone metadata does not match the Qwen-VL components"; + return nullptr; + } + + if (!is_oft && !is_fast) { + if (config.noise_seed >= 0) { + impl->noise_rng.seed(static_cast(config.noise_seed)); + } else { + std::random_device device; + std::seed_seq seed{device(), device(), device(), device(), + static_cast(Clock::now().time_since_epoch().count())}; + impl->noise_rng.seed(seed); + } + } + + if (config.verbosity >= 1) { + const char * variant_name = starvla_variant_name(impl->variant); + const char * policy_backend = is_oft ? impl->oft_policy->backend_name() + : (is_groot ? impl->groot_policy->backend_name() + : (is_pi ? impl->pi_policy->backend_name() + : (is_pi_v3 ? impl->pi_v3_policy->backend_name() + : impl->fast_policy->backend_name()))); + std::fprintf(stderr, + "%s: variant=%s policy=%s text=%s mmproj=%s profile=%s " + "qwen_backend=%s policy_backend=%s\n", + __func__, variant_name, impl->policy_path.string().c_str(), impl->text_path.string().c_str(), + impl->mmproj_path.string().c_str(), impl->normalization_profile_key.c_str(), + impl->qwen->backend_name(), policy_backend); + } + return std::unique_ptr(new StarVLAEngine(std::move(impl))); +} + +bool StarVLAEngine::predict(const observation & obs, StarVLAEngineResult & result, std::string & error) { + result = StarVLAEngineResult{}; + error.clear(); + const Clock::time_point total_start = Clock::now(); + const auto fail = [&]() { + result.actions.clear(); + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return false; + }; + + if (impl_ == nullptr || impl_->qwen == nullptr || + (!impl_->oft_policy && !impl_->groot_policy && !impl_->pi_policy && !impl_->pi_v3_policy && + !impl_->fast_policy)) { + error = "StarVLA engine is not initialized"; + return fail(); + } + + if (impl_->normalization == nullptr) { + error = "StarVLA normalization metadata is not initialized"; + return fail(); + } + std::string profile_error; + const NormalizationProfile * profile = + resolve_normalization_profile(*impl_->normalization, impl_->normalization_profile_key, profile_error); + if (profile == nullptr) { + error = "failed to select StarVLA normalization profile: " + profile_error; + return fail(); + } + std::vector normalized_actions; + std::string instruction; + + const auto make_noise = [&](size_t count, std::vector & noise) { + if (!obs.initial_noise.empty()) { + if (obs.initial_noise.size() != count || !std::all_of(obs.initial_noise.begin(), obs.initial_noise.end(), + [](float value) { return std::isfinite(value); })) { + error = "initial noise has an incompatible shape or non-finite value"; + return false; + } + noise = obs.initial_noise; + return true; + } + noise.resize(count); + std::normal_distribution normal(0.0f, 1.0f); + for (float & value : noise) { + value = ggml_bf16_to_fp32(ggml_fp32_to_bf16(normal(impl_->noise_rng))); + } + return true; + }; + + if (impl_->variant == StarVLAVariant::qwen25_fast) { + if (!obs.initial_noise.empty()) { + error = "StarVLA FAST does not use diffusion noise"; + return fail(); + } + const FastPolicyConfig & config = impl_->fast_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "FAST", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "FAST", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_fast_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA FAST prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + QwenVLGenerationConfig generation; + generation.max_length = config.generation_max_length; + generation.eos_token_ids = config.generation_eos_token_ids; + generation.top_k = config.generation_top_k; + generation.repetition_penalty = config.generation_repetition_penalty; + QwenVLGenerationResult generated; + stage_start = Clock::now(); + if (!impl_->qwen->generate_autoregressive(qwen_images, instruction, generation, generated, error)) { + error = "StarVLA FAST Qwen2.5-VL generation failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + if (generated.prompt_token_count == 0 || generated.full_sequence.size() < generated.prompt_token_count || + generated.full_sequence.size() > config.generation_max_length) { + error = "StarVLA FAST Qwen2.5-VL returned an incompatible generated sequence"; + return fail(); + } + stage_start = Clock::now(); + if (!impl_->fast_policy->decode_generated(generated.full_sequence, normalized_actions, error)) { + error = "StarVLA FAST codec decode failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->fast_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA FAST action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || normalized_actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA FAST returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen25_pi) { + const PIPolicyConfig & config = impl_->pi_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, true, "PI", error)) { + return fail(); + } + if (!obs.state.empty() && obs.state.size() != static_cast(config.state_dim)) { + error = + "StarVLA PI accepts either no state or exactly " + std::to_string(config.state_dim) + " state values"; + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> pre_resized_images; + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_pi_qwen_images(obs, config, pre_resized_images, processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_pi_v3_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA PI prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_layer_hidden_states(qwen_images, instruction, config.qwen_hidden_tuple_indices, + hidden_states, attention_mask, error)) { + error = "StarVLA PI Qwen2.5-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_hidden = static_cast(config.block_count) * attention_mask.size() * + static_cast(config.qwen_hidden_dim); + if (attention_mask.empty() || hidden_states.size() != expected_hidden) { + error = "StarVLA PI Qwen2.5-VL returned an incompatible layerwise conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->pi_policy->evaluate(hidden_states.data(), hidden_states.size(), + obs.state.empty() ? nullptr : obs.state.data(), obs.state.size(), noise.data(), + noise.size(), normalized_actions, error)) { + error = "StarVLA PI policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->pi_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA PI action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA PI returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen3_pi_v3) { + const PIV3PolicyConfig & config = impl_->pi_v3_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "PI_v3", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "PI_v3", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_pi_v3_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA PI_v3 prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_layer_hidden_states(qwen_images, instruction, config.qwen_hidden_tuple_indices, + hidden_states, attention_mask, error)) { + error = "StarVLA PI_v3 Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_hidden = + static_cast(config.qwen_layer_count) * attention_mask.size() * config.qwen_hidden_dim; + if (attention_mask.empty() || hidden_states.size() != expected_hidden) { + error = "StarVLA PI_v3 Qwen3-VL returned an incompatible layerwise conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->pi_v3_policy->evaluate(hidden_states.data(), hidden_states.size(), attention_mask.data(), + attention_mask.size(), noise.data(), noise.size(), normalized_actions, + error)) { + error = "StarVLA PI_v3 policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->pi_v3_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA PI_v3 action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA PI_v3 returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (impl_->variant == StarVLAVariant::qwen3_groot || impl_->variant == StarVLAVariant::qwen25_groot) { + const GR00TPolicyConfig & config = impl_->groot_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, false, "GR00T", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "GR00T", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_groot_instruction(config.cot_template, obs.task, instruction, error)) { + error = "failed to build the StarVLA GR00T prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector hidden_states; + std::vector attention_mask; + stage_start = Clock::now(); + if (!impl_->qwen->extract_full_hidden_states(qwen_images, instruction, hidden_states, attention_mask, error)) { + error = "StarVLA GR00T Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + if (hidden_states.empty() || hidden_states.size() % static_cast(config.qwen_hidden_dim) != 0 || + hidden_states.size() / static_cast(config.qwen_hidden_dim) != attention_mask.size()) { + error = "StarVLA GR00T Qwen3-VL returned an incompatible conditioning shape"; + return fail(); + } + + std::vector noise; + if (!make_noise(static_cast(config.horizon) * config.action_dim, noise)) { + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->groot_policy->evaluate(hidden_states.data(), hidden_states.size(), attention_mask.data(), + attention_mask.size(), noise.data(), noise.size(), normalized_actions, + error)) { + error = "StarVLA GR00T policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->groot_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA GR00T action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions || + !std::all_of(result.actions.begin(), result.actions.end(), + [](float action) { return std::isfinite(action); })) { + error = "StarVLA GR00T returned an incompatible or non-finite action tensor"; + return fail(); + } + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; + } + + if (!obs.initial_noise.empty()) { + error = "StarVLA OFT does not use diffusion noise"; + return fail(); + } + const OFTPolicyConfig & config = impl_->oft_policy->config(); + if (!validate_observation(obs, config.image_count, config.image_names, true, "OFT", error)) { + return fail(); + } + + Clock::time_point stage_start = Clock::now(); + std::vector> processed_images; + std::vector qwen_images; + if (!prepare_qwen_images(obs, config, "OFT", processed_images, qwen_images, error)) { + return fail(); + } + result.timings.image_preprocess_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!build_oft_instruction(config.prompt, obs.task, obs.state, instruction, error)) { + error = "failed to build the StarVLA OFT prompt: " + error; + return fail(); + } + result.timings.prompt_ms = elapsed_ms(stage_start, Clock::now()); + + std::vector action_queries; + stage_start = Clock::now(); + if (!impl_->qwen->extract_token_embeddings(qwen_images, instruction, config.action_token_id, + static_cast(config.horizon), action_queries, error)) { + error = "StarVLA OFT Qwen3-VL inference failed: " + error; + return fail(); + } + result.timings.qwen3vl_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_queries = static_cast(config.horizon) * config.input_dim; + if (action_queries.size() != expected_queries) { + error = "StarVLA OFT Qwen3-VL returned an incompatible action-query shape"; + return fail(); + } + + stage_start = Clock::now(); + if (!impl_->oft_policy->evaluate(action_queries.data(), action_queries.size(), normalized_actions, error)) { + error = "StarVLA OFT policy inference failed: " + error; + return fail(); + } + result.timings.policy_ms = elapsed_ms(stage_start, Clock::now()); + + stage_start = Clock::now(); + if (!impl_->oft_policy->unnormalize(normalized_actions, profile->key, result.actions, error)) { + error = "StarVLA OFT action unnormalization failed: " + error; + return fail(); + } + result.timings.unnormalize_ms = elapsed_ms(stage_start, Clock::now()); + const size_t expected_actions = static_cast(config.horizon) * config.action_dim; + if (result.actions.size() != expected_actions) { + error = "StarVLA OFT returned an incompatible action tensor shape"; + return fail(); + } + for (float action : result.actions) { + if (!std::isfinite(action)) { + error = "StarVLA OFT returned a non-finite unnormalized action"; + return fail(); + } + } + + result.chunk_size = config.horizon; + result.action_dim = config.action_dim; + result.timings.total_ms = elapsed_ms(total_start, Clock::now()); + return true; +} + +void StarVLAEngine::reset() { + if (impl_ != nullptr && impl_->qwen != nullptr) { + impl_->qwen->reset(); + } +} + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_engine.h b/src/models/starvla/starvla_engine.h new file mode 100644 index 0000000..92c7319 --- /dev/null +++ b/src/models/starvla/starvla_engine.h @@ -0,0 +1,74 @@ +#pragma once + +#include "models/model.h" + +#include +#include +#include +#include + +namespace robotcpp::starvla { + +enum class StarVLAVariant { + qwen3_oft, + qwen3_groot, + qwen3_pi_v3, + qwen25_oft, + qwen25_groot, + qwen25_pi, + qwen25_fast, +}; + +const char * starvla_variant_name(StarVLAVariant variant) noexcept; +const char * starvla_variant_framework(StarVLAVariant variant) noexcept; +bool starvla_variant_from_metadata(const std::string & framework, const std::string & backbone, + StarVLAVariant & variant) noexcept; + +struct StarVLAEngineConfig { + std::string policy_path; + std::string text_path; + std::string mmproj_path; + int n_threads = 0; + int n_ctx = 2048; + int n_batch = 512; + int64_t noise_seed = -1; + int verbosity = 0; +}; + +struct StarVLAStageTimings { + double image_preprocess_ms = 0.0; + double prompt_ms = 0.0; + double qwen3vl_ms = 0.0; + double policy_ms = 0.0; + double unnormalize_ms = 0.0; + double total_ms = 0.0; +}; + +struct StarVLAEngineResult { + std::vector actions; + int chunk_size = 0; + int action_dim = 0; + StarVLAStageTimings timings; +}; + +class StarVLAEngine { + public: + ~StarVLAEngine(); + + StarVLAEngine(const StarVLAEngine &) = delete; + StarVLAEngine & operator=(const StarVLAEngine &) = delete; + + static std::unique_ptr load(const StarVLAEngineConfig & config, std::string & error); + + bool predict(const observation & obs, StarVLAEngineResult & result, std::string & error); + void reset(); + + private: + struct Impl; + + explicit StarVLAEngine(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +} // namespace robotcpp::starvla diff --git a/src/models/starvla/starvla_model.cpp b/src/models/starvla/starvla_model.cpp new file mode 100644 index 0000000..e320676 --- /dev/null +++ b/src/models/starvla/starvla_model.cpp @@ -0,0 +1,93 @@ +#include "models/starvla/starvla_model.h" + +#include "models/starvla/starvla_engine.h" + +#include +#include +#include + +namespace robotcpp { +namespace { + +void add_metric(model_result & out, const char * name, double value) { + model_metric metric; + metric.name = name; + metric.value = value; + out.metrics.push_back(std::move(metric)); +} + +} // namespace + +StarVLAModel::StarVLAModel(std::unique_ptr engine) : engine_(std::move(engine)) {} + +StarVLAModel::~StarVLAModel() = default; + +const char * StarVLAModel::type() const { + return "starvla"; +} + +bool StarVLAModel::predict(const observation & obs, model_result & out, std::string & error) { + out = model_result{}; + error.clear(); + if (engine_ == nullptr) { + error = "StarVLA model is not initialized"; + return false; + } + + starvla::StarVLAEngineResult result; + if (!engine_->predict(obs, result, error)) { + return false; + } + out.actions = std::move(result.actions); + out.chunk_size = result.chunk_size; + out.action_dim = result.action_dim; + add_metric(out, "image_preprocess_ms", result.timings.image_preprocess_ms); + add_metric(out, "prompt_ms", result.timings.prompt_ms); + add_metric(out, "qwen3vl_ms", result.timings.qwen3vl_ms); + add_metric(out, "policy_ms", result.timings.policy_ms); + add_metric(out, "unnormalize_ms", result.timings.unnormalize_ms); + add_metric(out, "model_total_ms", result.timings.total_ms); + return true; +} + +void StarVLAModel::reset() { + if (engine_ != nullptr) { + engine_->reset(); + } +} + +bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error) { + out.reset(); + error.clear(); + if (!is_starvla_model_type(args.type)) { + error = std::string("model type '") + model_type_name(args.type) + "' is not a StarVLA model type"; + return false; + } + if (args.noise_mode != 0) { + error = "StarVLA does not support SmolVLA --noise-mode debug-sin; use Gaussian noise and --noise-seed"; + return false; + } + starvla::StarVLAEngineConfig config; + config.policy_path = args.policy_path; + config.text_path = args.llm_path; + config.mmproj_path = args.mmproj_path; + config.n_threads = args.threads; + config.n_ctx = args.n_ctx; + config.n_batch = args.n_batch; + config.noise_seed = args.noise_seed; + config.verbosity = args.verbosity; + std::unique_ptr engine = starvla::StarVLAEngine::load(config, error); + if (engine == nullptr) { + return false; + } + + std::unique_ptr model(new (std::nothrow) StarVLAModel(std::move(engine))); + if (model == nullptr) { + error = "failed to allocate StarVLA model"; + return false; + } + out = std::move(model); + return true; +} + +} // namespace robotcpp diff --git a/src/models/starvla/starvla_model.h b/src/models/starvla/starvla_model.h new file mode 100644 index 0000000..9c41099 --- /dev/null +++ b/src/models/starvla/starvla_model.h @@ -0,0 +1,35 @@ +#pragma once + +#include "models/model.h" + +#include +#include + +namespace robotcpp::starvla { +class StarVLAEngine; +} + +namespace robotcpp { + +class StarVLAModel final : public Model { + public: + ~StarVLAModel() override; + + StarVLAModel(const StarVLAModel &) = delete; + StarVLAModel & operator=(const StarVLAModel &) = delete; + + const char * type() const override; + bool predict(const observation & obs, model_result & out, std::string & error) override; + void reset() override; + + private: + explicit StarVLAModel(std::unique_ptr engine); + + friend bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error); + + std::unique_ptr engine_; +}; + +bool make_starvla_model(const model_args & args, std::unique_ptr & out, std::string & error); + +} // namespace robotcpp diff --git a/tests/starvla/model_test.cpp b/tests/starvla/model_test.cpp new file mode 100644 index 0000000..73f665b --- /dev/null +++ b/tests/starvla/model_test.cpp @@ -0,0 +1,59 @@ +#include "models/model.h" +#include "models/starvla/starvla_engine.h" + +#include +#include +#include +#include + +namespace { + +using robotcpp::starvla::StarVLAVariant; + +struct VariantCase { + const char * framework; + const char * backbone; + const char * name; + StarVLAVariant variant; +}; + +constexpr std::array kVariants = {{ + // Qwen3-VL + {"oft", "qwen3_vl", "qwen3_oft", StarVLAVariant::qwen3_oft}, + {"groot", "qwen3_vl", "qwen3_groot", StarVLAVariant::qwen3_groot}, + {"pi_v3", "qwen3_vl", "qwen3_pi_v3", StarVLAVariant::qwen3_pi_v3}, + + // Qwen2.5-VL + {"oft", "qwen2_5_vl", "qwen25_oft", StarVLAVariant::qwen25_oft}, + {"groot", "qwen2_5_vl", "qwen25_groot", StarVLAVariant::qwen25_groot}, + {"pi", "qwen2_5_vl", "qwen25_pi", StarVLAVariant::qwen25_pi}, + {"fast", "qwen2_5_vl", "qwen25_fast", StarVLAVariant::qwen25_fast}, +}}; + +} // namespace + +int main() { + for (const VariantCase & test : kVariants) { + StarVLAVariant variant = StarVLAVariant::qwen3_oft; + if (!robotcpp::starvla::starvla_variant_from_metadata( + test.framework, test.backbone, variant) || + variant != test.variant || + std::string(robotcpp::starvla::starvla_variant_name(variant)) != test.name || + std::string(robotcpp::starvla::starvla_variant_framework(variant)) != + test.framework) { + std::fprintf(stderr, "variant check failed: %s\n", test.name); + return 1; + } + } + + robotcpp::model_args args; + args.type = robotcpp::model_type::starvla; + std::unique_ptr model; + std::string error; + if (robotcpp::make_model(args, model, error) || model || + error.find("policy path is required") == std::string::npos) { + std::fprintf(stderr, "unexpected factory result: %s\n", error.c_str()); + return 1; + } + return 0; +} diff --git a/tools/apply_patches.sh b/tools/apply_patches.sh new file mode 100755 index 0000000..746584c --- /dev/null +++ b/tools/apply_patches.sh @@ -0,0 +1,115 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly EXPECTED_LLAMA_COMMIT="3e941b813b1acbbf06c2203a94ceb33d84748c1e" +readonly SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +readonly REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)" +readonly LLAMA_DIR="${LLAMA_CPP_DIR:-${REPO_ROOT}/third_party/llama.cpp}" +readonly PATCH_DIR="${REPO_ROOT}/patches/llama.cpp" +readonly PATCHES=( + "${PATCH_DIR}/0001-qwen3vl-vision-parity.patch" + "${PATCH_DIR}/0002-per-context-native-graph-control.patch" +) + +usage() { + cat <<'EOF' +Usage: tools/apply_patches.sh [--check|--revert] + +With no option, apply the repository patches to third-party checkouts. + --check Validate pinned revisions and report patch state. + --revert Remove an already applied complete patch set. + +Set LLAMA_CPP_DIR to validate or patch another checkout of the pinned revision. +EOF +} + +mode="apply" +case "${1:-}" in + "") ;; + --check) mode="check" ;; + --revert) mode="revert" ;; + -h|--help) usage; exit 0 ;; + *) usage >&2; exit 2 ;; +esac + +if [[ ! -d "${LLAMA_DIR}/.git" && ! -f "${LLAMA_DIR}/.git" ]]; then + echo "error: llama.cpp checkout not found: ${LLAMA_DIR}" >&2 + exit 1 +fi + +actual_commit="$(git -C "${LLAMA_DIR}" rev-parse HEAD)" +if [[ "${actual_commit}" != "${EXPECTED_LLAMA_COMMIT}" ]]; then + echo "error: unsupported llama.cpp revision" >&2 + echo " expected: ${EXPECTED_LLAMA_COMMIT}" >&2 + echo " actual: ${actual_commit}" >&2 + exit 1 +fi + +states=() +for patch in "${PATCHES[@]}"; do + if git -C "${LLAMA_DIR}" apply --reverse --check "${patch}" >/dev/null 2>&1; then + states+=("applied") + elif git -C "${LLAMA_DIR}" apply --check "${patch}" >/dev/null 2>&1; then + states+=("pending") + else + echo "error: patch is neither cleanly applicable nor already applied: ${patch}" >&2 + exit 1 + fi +done + +all_pending=true +all_applied=true +for state in "${states[@]}"; do + [[ "${state}" == "pending" ]] || all_pending=false + [[ "${state}" == "applied" ]] || all_applied=false +done + +if [[ "${mode}" == "check" ]]; then + for i in "${!PATCHES[@]}"; do + printf '%-8s %s\n' "${states[$i]}" "${PATCHES[$i]#${REPO_ROOT}/}" + done + if ! ${all_pending} && ! ${all_applied}; then + echo "error: partial patch set detected" >&2 + exit 1 + fi + exit 0 +fi + +if [[ "${mode}" == "apply" ]]; then + if ${all_applied}; then + echo "Repository patches are already applied." + exit 0 + fi + if ! ${all_pending}; then + echo "error: refusing to apply a partial patch set" >&2 + exit 1 + fi + if [[ -n "$(git -C "${LLAMA_DIR}" status --porcelain)" ]]; then + echo "error: refusing to patch a dirty llama.cpp checkout" >&2 + exit 1 + fi + for patch in "${PATCHES[@]}"; do + git -C "${LLAMA_DIR}" apply --check "${patch}" + done + for patch in "${PATCHES[@]}"; do + git -C "${LLAMA_DIR}" apply "${patch}" + echo "applied ${patch#${REPO_ROOT}/}" + done + exit 0 +fi + +if ${all_pending}; then + echo "Repository patches are not applied." + exit 0 +fi +if ! ${all_applied}; then + echo "error: refusing to revert a partial patch set" >&2 + exit 1 +fi +for ((i=${#PATCHES[@]} - 1; i >= 0; --i)); do + git -C "${LLAMA_DIR}" apply --reverse --check "${PATCHES[$i]}" +done +for ((i=${#PATCHES[@]} - 1; i >= 0; --i)); do + git -C "${LLAMA_DIR}" apply --reverse "${PATCHES[$i]}" + echo "reverted ${PATCHES[$i]#${REPO_ROOT}/}" +done diff --git a/tools/hf2gguf/README.md b/tools/hf2gguf/README.md index 57575c2..2f4260c 100644 --- a/tools/hf2gguf/README.md +++ b/tools/hf2gguf/README.md @@ -8,6 +8,8 @@ This directory contains tools for converting checkpoints to GGUF. - `smolvla/`: converts LeRobot-style SmolVLA checkpoints into four GGUF components. - `pi0/`: converts LeRobot-style pi0 checkpoints into six split GGUF components. +- `starvla/`: converts seven StarVLA Qwen3-VL and Qwen2.5-VL checkpoints. See + [`starvla/README.md`](starvla/README.md) for supported variants and commands. - `environment.yaml`: conda environment for the converters. ## Usage diff --git a/tools/hf2gguf/README_ZH.md b/tools/hf2gguf/README_ZH.md index 4219bff..ad1cec1 100644 --- a/tools/hf2gguf/README_ZH.md +++ b/tools/hf2gguf/README_ZH.md @@ -8,6 +8,8 @@ - `smolvla/`：将 SmolVLA的lerobot-style的checkpoint 转成四个 GGUF component。 - `pi0/`：将 pi0的lerobot-style的checkpoint 转成六个 split GGUF component。 +- `starvla/`：转换七个 StarVLA Qwen3-VL 和 Qwen2.5-VL checkpoint，支持范围和命令见 + [`starvla/README.md`](starvla/README.md)。 - `environment.yaml`：converter conda 环境。 ## 使用说明 diff --git a/tools/hf2gguf/environment.yaml b/tools/hf2gguf/environment.yaml index a27464c..f014d70 100644 --- a/tools/hf2gguf/environment.yaml +++ b/tools/hf2gguf/environment.yaml @@ -8,6 +8,8 @@ dependencies: - torch - numpy - safetensors + - huggingface_hub>=0.36.0 - sentencepiece - transformers==4.53.2 + - pillow==12.1.1 - pyyaml diff --git a/tools/hf2gguf/starvla/README.md b/tools/hf2gguf/starvla/README.md new file mode 100644 index 0000000..a663e35 --- /dev/null +++ b/tools/hf2gguf/starvla/README.md @@ -0,0 +1,177 @@ +# Converting StarVLA checkpoints + +The scripts in this directory download StarVLA checkpoints and convert them to +GGUF files used by robot.cpp. + +## Models + +| Variant | Backbone | Policy | +| --- | --- | --- | +| `oft` | Qwen3-VL | OFT | +| `groot` | Qwen3-VL | GR00T | +| `pi_v3` | Qwen3-VL | PI_v3 | +| `qwen25_oft` | Qwen2.5-VL | OFT | +| `qwen25_groot` | Qwen2.5-VL | GR00T | +| `qwen25_pi` | Qwen2.5-VL | PI | +| `qwen25_fast` | Qwen2.5-VL | FAST | + +All variants use `starvla` as the public model type. The loader reads the +backbone and policy type from the policy GGUF. + +The first six variants contain a BF16 Qwen file, a BF16 multimodal projector, +and an FP32 policy file. FAST uses the fine-tuned BF16 Qwen model as its policy; +its separate policy GGUF contains the integer token map and codec data. The +loader checks the bundle UUID to prevent files from different conversions from +being combined. + +Qwen3 FAST is not listed because StarVLA has not published a fine-tuned Qwen3 +FAST policy checkpoint. + +[`checkpoint_catalog.json`](checkpoint_catalog.json) defines the supported +topologies and pins the official release files and shared Qwen assets. + +## Environment + +```bash +conda env create -f tools/hf2gguf/starvla/environment.yaml +conda activate starvla_gguf_converter +``` + +The scripts use `.venv/bin/python` by default. Set `PYTHON=python` to use the +active conda environment. + +## Convert an official release + +Convert one of the variants from the table above: + +```bash +tools/hf2gguf/starvla/convert.sh oft +``` + +This downloads and verifies the catalog checkpoint, prepares a +clean llama.cpp worktree at the pinned revision, converts all components, and +validates the resulting bundle. It refuses to overwrite an existing output +directory. Pass a second argument to select another output directory: + +```bash +tools/hf2gguf/starvla/convert.sh qwen25_fast /path/to/output +``` + +## Convert a training checkpoint + +Current StarVLA training runs contain the files needed by the converter: + +```text +/ + config.yaml + dataset_statistics.json + checkpoints/steps__pytorch_model.pt + # or checkpoints/steps__model.safetensors +``` + +Pass the checkpoint and the matching topology from the model table: + +```bash +tools/hf2gguf/starvla/convert.sh oft /path/to/output \ + --checkpoint /path/to/run/checkpoints/steps_5000_model.safetensors +``` + +The run directory is inferred from checkpoints under `checkpoints/` or +`final_model/`. Use `--source-dir /path/to/run` when the files use another +layout. If `dataset_statistics.json` has several profiles and does not contain +the catalog default, select one with `--unnorm-key`: + +```bash +tools/hf2gguf/starvla/convert.sh groot /path/to/output \ + --checkpoint /path/to/run/final_model/pytorch_model.pt \ + --unnorm-key bridge_dataset +``` + +Supported training exports are flat PyTorch state dictionaries (`.pt`) and +flat safetensors files (`.safetensors`) written by `train_starvla.py`, including +periodic and final checkpoints. The converter does not consume optimizer +state, distributed checkpoint shards, or a checkpoint whose architecture no +longer matches the selected variant. `config.json` and `config.full.yaml` are +not required. + +The converter hashes the local checkpoint and run metadata, so its bundle UUID +and manifest differ from the official release even when the weights are equal. + +A successful conversion writes exactly four files: + +```text +qwen--bf16.gguf +mmproj--bf16.gguf +starvla--policy-fp32.gguf +conversion_manifest.json +``` + +For FAST, the policy GGUF stores the integer token map and codec data instead +of FP32 policy weights. Its filenames are: + +```text +qwen-qwen25-fast-bf16.gguf +mmproj-qwen25-fast-bf16.gguf +policy-qwen25-fast.gguf +conversion_manifest.json +``` + +Set `STARVLA_LOCAL_FILES_ONLY=1` to forbid network access and use already +downloaded sources. The low-level converters remain available for debugging, +but normal conversion should use `convert.sh` so all paths and revisions come +from [`checkpoint_catalog.json`](checkpoint_catalog.json). + +## Build + +The runtime needs two llama.cpp patches maintained in this repository. See +[`patches/llama.cpp/README.md`](../../../patches/llama.cpp/README.md) for their +scope. + +```bash +./tools/apply_patches.sh +cmake -S . -B build_cuda \ + -DGGML_CUDA=ON \ + -DBUILD_TESTING=ON \ + -DROBOT_CPP_BUILD_STARVLA=ON \ + -DROBOT_CPP_BUILD_MODEL_CLI=ON +cmake --build build_cuda -j +``` + +## Run + +Pass the Qwen, multimodal projector, and policy GGUF files separately: + +```bash +CUDA_VISIBLE_DEVICES=0 build_cuda/bin/model-cli \ + --model-type starvla \ + --policy ckpts/starvla/gguf/oft/starvla-oft-policy-fp32.gguf \ + --llm ckpts/starvla/gguf/oft/qwen-oft-bf16.gguf \ + --mmproj ckpts/starvla/gguf/oft/mmproj-oft-bf16.gguf \ + --image /path/to/frame-224-rgb.png \ + --image-name image_0 \ + --task "grab the block." \ + --n-ctx 2048 \ + --n-batch 2048 +``` + +GR00T, PI_v3, and PI accept `--noise-seed`. FAST accepts one RGB `image_0` and +no robot state. + +The server uses the same model type and policy file: + +```bash +CUDA_VISIBLE_DEVICES=0 build_cuda/bin/model-server \ + --model-type starvla \ + --policy ckpts/starvla/gguf/oft/starvla-oft-policy-fp32.gguf \ + --llm ckpts/starvla/gguf/oft/qwen-oft-bf16.gguf \ + --mmproj ckpts/starvla/gguf/oft/mmproj-oft-bf16.gguf \ + --host 127.0.0.1 \ + --port 5555 \ + --n-ctx 2048 \ + --n-batch 2048 +``` + +The policy GGUF records its default action normalization profile. + +The repository does not include upstream checkpoints. Check each model's +license before distributing converted files. diff --git a/tools/hf2gguf/starvla/__init__.py b/tools/hf2gguf/starvla/__init__.py new file mode 100755 index 0000000..52b10f1 --- /dev/null +++ b/tools/hf2gguf/starvla/__init__.py @@ -0,0 +1 @@ +"""StarVLA checkpoint conversion tools.""" diff --git a/tools/hf2gguf/starvla/checkpoint_catalog.json b/tools/hf2gguf/starvla/checkpoint_catalog.json new file mode 100644 index 0000000..040c010 --- /dev/null +++ b/tools/hf2gguf/starvla/checkpoint_catalog.json @@ -0,0 +1,478 @@ +{ + "schema_version": 1, + "source_revisions": { + "starvla": "631aae02afe6d95876e923ff518e8ff2ab9a2f88", + "llama_cpp": "3e941b813b1acbbf06c2203a94ceb33d84748c1e" + }, + "shared_assets": { + "qwen3_vl_4b_instruct": { + "directory": "qwen3-vl-4b-instruct", + "repo_id": "Qwen/Qwen3-VL-4B-Instruct", + "revision": "ebb281ec70b05090aa6165b016eac8ec08e71b17", + "files": [ + "chat_template.json", + "config.json", + "generation_config.json", + "merges.txt", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", + "vocab.json" + ], + "file_hashes": { + "chat_template.json": {"size": 5502, "sha256": "6f8a6a55027e3da5160105556cda5dd69f6423f1c32645f6730d32de7773d0c4"}, + "config.json": {"size": 1505, "sha256": "edac7703329133edfc53e46ac0081835144c99d7eebf28b71c732694d435224d"}, + "generation_config.json": {"size": 269, "sha256": "8469742d1fce0de951c8909b26a2c0c0d8490837ce476efb114da9e0cefc4d44"}, + "merges.txt": {"size": 1671839, "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3"}, + "preprocessor_config.json": {"size": 390, "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516"}, + "tokenizer.json": {"size": 7032403, "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7"}, + "tokenizer_config.json": {"size": 10868, "sha256": "c2da771801886ad9ae98181793ffd3dfb7f1af30f6f7c6a4e15d7dbba52e2399"}, + "video_preprocessor_config.json": {"size": 385, "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 1507, "sha256": "ef6ec5fd4c5a80b549208f2352d88c480436db2cf9983359f23260c36e4ae38d"} + } + }, + "qwen2_5_vl_3b_instruct": { + "directory": "qwen2.5-vl-3b-instruct", + "repo_id": "Qwen/Qwen2.5-VL-3B-Instruct", + "revision": "66285546d2b821cf421d4f5eb2576359d3770cd3", + "files": [ + "chat_template.json", + "config.json", + "generation_config.json", + "merges.txt", + "model.safetensors.index.json", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer_config.json", + "vocab.json" + ], + "file_hashes": { + "chat_template.json": {"size": 1050, "sha256": "ad60d90252ed0b0705ba14e2d0ad0fec0beac1ea955642b54059b36052d8bc96"}, + "config.json": {"size": 1373, "sha256": "7ed3eed5be6924cc800e8a5e53fc405c1aab1aaf36bad65c33403b36c56827f5"}, + "generation_config.json": {"size": 216, "sha256": "533f191cc257b7de37a4fccd0a7a1706d75e1aa660f93efaa54e5a2a9f9aace9"}, + "merges.txt": {"size": 1671839, "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3"}, + "model.safetensors.index.json": {"size": 65448, "sha256": "c7dd78a4c6bea60b51332f1baf37b8f8124ecab2c35395a29a29825bf2619768"}, + "preprocessor_config.json": {"size": 350, "sha256": "f2058c716eef96ccaed1cc1e2d0c08306b62586d535b28d9d08e691b2fab7ca0"}, + "tokenizer.json": {"size": 7031645, "sha256": "c0382117ea329cdf097041132f6d735924b697924d6f6fc3945713e96ce87539"}, + "tokenizer_config.json": {"size": 5702, "sha256": "4abd3520120e266da84c0864fee064d1fb10806f02225911a47253dd38dc5f56"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 1375, "sha256": "9c22fba5261a8e47aa66be0e4ef22473190168859dc3bbe7f283fbc4f161b0eb"} + }, + "optional_weight_files": [ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors" + ], + "optional_weight_hashes": { + "model-00001-of-00002.safetensors": {"size": 3982649232, "sha256": "41a8895c164b4d32bae6b302f4603fcbc1797f32dafa45c7e9bcda23c6755df8"}, + "model-00002-of-00002.safetensors": {"size": 3526688744, "sha256": "365531ff8752420e89dee707b79d021fb2d6e25abafe486f080555a4fe6972e4"} + } + }, + "qwen2_5_vl_3b_instruct_action": { + "directory": "qwen2.5-vl-3b-instruct-action", + "repo_id": "StarVLA/Qwen2.5-VL-3B-Instruct-Action", + "revision": "ce86bd9a53416527b8361e8dfc47316288ffa110", + "files": [ + "README.md", + "added_token_id_map.json", + "added_tokens.json", + "chat_template.jinja", + "config.json", + "generation_config.json", + "merges.txt", + "model.safetensors.index.json", + "preprocessor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json", + "video_preprocessor_config.json", + "vocab.json" + ], + "file_hashes": { + "README.md": {"size": 482, "sha256": "cef1e9f3b90d50d1b6274fed603140127de9441bd6d8425811225d00a290f46f"}, + "added_token_id_map.json": {"size": 66476, "sha256": "a774a771870979578111a9f083e03e421bc3e6c0c7070d81e775acc21c74a21a"}, + "added_tokens.json": {"size": 67079, "sha256": "fcca65c62c6da071c4046abbc18b8287c51030e62febec408bac498a03652eaa"}, + "chat_template.jinja": {"size": 1017, "sha256": "a0bc6f6fc7a29a80017a433e8f03a1cc1236e838a944a2d034295a60c4f2fddb"}, + "config.json": {"size": 3317, "sha256": "5c30acf44442bbdd863b87a6f61b6879616a2933271bf62841de14037f6c0f7d"}, + "generation_config.json": {"size": 244, "sha256": "76001fd927297f839d96c5a52dd09de3406a4c28822fba4f525d13c5a1e2c8d7"}, + "merges.txt": {"size": 1671853, "sha256": "8831e4f1a044471340f7c0a83d7bd71306a5b867e95fd870f74d0c5308a904d5"}, + "model.safetensors.index.json": {"size": 65484, "sha256": "e6ce326cff552529deb7ca2e05616b9b79fba9d26632309564744c3731fbb644"}, + "preprocessor_config.json": {"size": 829, "sha256": "dfc7263fe735989c65c290d394198c4724d5afc58bf15c815d5e2e25b00b51b1"}, + "special_tokens_map.json": {"size": 312563, "sha256": "707f14d06c06e20212dbe5c118873f1c024f28ffac4b07c52bd2840ba0c34290"}, + "tokenizer.json": {"size": 11822194, "sha256": "07da2a694acc4f6e63d67da9926817ee35b0354b1e570a6a73d325760a1c2ed2"}, + "tokenizer_config.json": {"size": 438450, "sha256": "7cd59c7a865d2989c0d9b18bf485a5ee212f8ff333acd305a1bf560b75c16575"}, + "video_preprocessor_config.json": {"size": 913, "sha256": "15bb7c2f2bc95fe9cc3749a4b287872b4886a15e9fcf550a4122a46ae26150bd"}, + "vocab.json": {"size": 2776833, "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910"} + }, + "staged_overrides": { + "config.json": {"size": 3350, "sha256": "782edd73d2c9584d65350a6410780b96bef658437cbd9d8e0ed7006a1e3fcaed"} + }, + "optional_weight_files": [ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors" + ], + "optional_weight_hashes": { + "model-00001-of-00002.safetensors": {"size": 4959940464, "sha256": "0abe459fc004959698441fd706b7721ec4633e963284b1b260df31ed5f765960"}, + "model-00002-of-00002.safetensors": {"size": 2556676080, "sha256": "a7a84f03ce697eddedc3abba904e2a110104fe78d6be9a0067f5b8f45c358e1c"} + } + }, + "fast_codec": { + "directory": "fast-codec", + "repo_id": "physical-intelligence/fast", + "revision": "ec4d7aa71691cac0b8bed6942be45684db2110f4", + "files": [ + "processing_action_tokenizer.py", + "processor_config.json", + "special_tokens_map.json", + "tokenizer.json", + "tokenizer_config.json" + ], + "file_hashes": { + "processing_action_tokenizer.py": {"size": 6145, "sha256": "6f021ca1f4c1b194ab6fa399d80baf3d642eadb17efb8f73301e4ac401522c20"}, + "processor_config.json": {"size": 253, "sha256": "f40cfbb1020858fe1d48c0f946b0c1315a90d6e84aa82710036f24f4c167706a"}, + "special_tokens_map.json": {"size": 3, "sha256": "ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356"}, + "tokenizer.json": {"size": 686974, "sha256": "6507dd709287fd018882120c0071787f1f62bad9f180f1e8c5235bda1b71fa78"}, + "tokenizer_config.json": {"size": 322, "sha256": "b4030e2a13a0dea22e99d54c086fb320c71e66ad034ac4eba4301a0a27d5e5cd"} + } + } + }, + "variants": { + "oft": { + "model_type": "starvla", + "framework": "oft", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "oft-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-OFT-Bridge-RT-1", + "revision": "c3fc8f028429ba14819bf3b16e098776b670c889", + "files": [ + "config.json", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.json": {"size": 3920, "sha256": "6a6b0dd11ec26f88aca711a8886ecb619bbb895846a6353df53f04ece682b318"}, + "config.yaml": {"size": 3207, "sha256": "6c074974697115284b1624dda5230f3dd27e1d9a373db73006467004ed859c2a"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"} + }, + "checkpoint": { + "path": "checkpoints/steps_5000_pytorch_model.pt", + "size": 9785060316, + "sha256": "371cb744227687bb99bcad7f9ff2250cf06da75631359ad3eba4c6bc52570607" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 730, + "vlm_tensors": 714, + "policy_tensors": 16, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 4892395015, + "vlm_numel": 4826771968, + "policy_numel": 65623047, + "dtypes": { + "bfloat16": 730 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560], + "action_model.model.fc1.weight": [5120, 2560], + "action_model.model.fc2.weight": [7, 5120] + } + }, + "groot": { + "model_type": "starvla", + "framework": "groot", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "groot-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-GR00T-Bridge-RT-1", + "revision": "12acc0b0f1f6230df21c479934a67a930b52f878", + "files": [ + "config.json", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.json": {"size": 3926, "sha256": "9efddc3c21039fa473823080a939dfa686050a8a2d4a4cb1b01b1a06913fccf5"}, + "config.yaml": {"size": 3174, "sha256": "01e092e9a3a9380885f1a27953048e7ed1ef7f6c99ea1984d289d1358e4ba85f"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"} + }, + "checkpoint": { + "path": "checkpoints/steps_20000_pytorch_model.pt", + "size": 9976845210, + "sha256": "769d6c400d582a86ae8df8b0b445240ab679dbe77eeb72a4db71e43cd129c7c3" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 962, + "vlm_tensors": 714, + "policy_tensors": 248, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 4988244743, + "vlm_numel": 4826771968, + "policy_numel": 161472775, + "dtypes": { + "bfloat16": 962 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560] + } + }, + "pi_v3": { + "model_type": "starvla", + "framework": "pi_v3", + "backbone": "qwen3_vl", + "qwen_asset": "qwen3_vl_4b_instruct", + "default_unnorm_key": "oxe_bridge", + "directory": "pi-v3-bridge-rt1", + "repo_id": "StarVLA/Qwen3VL-PI_v3-Bridge-RT_1", + "revision": "99a3c01b3977e6442871a1fb62ce178279c5c3ed", + "files": [ + "config.full.yaml", + "config.yaml", + "dataset_statistics.json" + ], + "file_hashes": { + "config.full.yaml": {"size": 3087, "sha256": "85ff9fba2c9426d35f12efabfd52b9ef4633d91d084421952b391809ea73b33f"}, + "config.yaml": {"size": 1915, "sha256": "f266bd2de5b9fb7078c8314954e98f72476a36a8944e3f380096ba5b1256901b"}, + "dataset_statistics.json": {"size": 5987, "sha256": "9925e884e37ca807061b5d41206bcf72814300e7effb5b4edc067dc821fca004"} + }, + "checkpoint": { + "path": "checkpoints/steps_50000_pytorch_model.pt", + "size": 10922634912, + "sha256": "7f59a5d0fa9c167fabd941bca8e606bdf5597bfb4f99ca83e345672dd9c345ed" + }, + "policy_prefixes": [ + "action_model.", + "project_layers." + ], + "expected": { + "total_tensors": 1386, + "vlm_tensors": 714, + "policy_tensors": 672, + "visual_tensors": 315, + "text_tensors": 398, + "lm_head_tensors": 1, + "total_numel": 5461066247, + "vlm_numel": 4826771968, + "policy_numel": 634294279, + "dtypes": { + "bfloat16": 1386 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.language_model.embed_tokens.weight": [151936, 2560], + "lm_head.weight": [151936, 2560] + } + }, + "qwen25_oft": { + "model_type": "starvla", + "framework": "oft", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct", + "default_unnorm_key": "bridge_dataset", + "directory": "qwen25-oft-bridge-rt1", + "repo_id": "StarVLA/Qwen-OFT-Bridge-RT-1", + "revision": "11fa6440835ba3e912de43cfe8521043360ffc02", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 2876, "sha256": "2f0362a1c0ae1eafba90d0feadf34652515c1d8a0c956db8f97e532492f2cdab"}, + "dataset_statistics.json": {"size": 6007, "sha256": "d2c4803c94d3b6eb1b8e8e100280e16a53b6058c8c2a7e747d27ecf9fcf9a4de"}, + "summary.jsonl": {"size": 33, "sha256": "a352646601877394e54b68ce09697098866d5904b825d7932826decaef0b2f8f"} + }, + "checkpoint": { + "path": "checkpoints/steps_10000_pytorch_model.pt", + "size": 8215912766, + "sha256": "51fe8d22c8d57116c2f59c5fdb24323fa3411149e888b807edba99b8354e0861" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 841, + "vlm_tensors": 825, + "policy_tensors": 16, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4107800583, + "vlm_numel": 4065787904, + "policy_numel": 42012679, + "total_nbytes": 8215601166, + "vlm_nbytes": 8131575808, + "policy_nbytes": 84025358, + "dtypes": { + "bfloat16": 841 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [151936, 2048], + "lm_head.weight": [151936, 2048] + } + }, + "qwen25_groot": { + "model_type": "starvla", + "framework": "groot", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "oxe_bridge", + "directory": "qwen25-groot-bridge-rt1", + "repo_id": "StarVLA/Qwen-GR00T-Bridge-RT-1", + "revision": "5ebc661ba38b29c28f20fff6574801e6f49f3466", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 3175, "sha256": "80d36dd087bd8d0feff246be94a7edcb296161823bce6340f76cc253724fbf1d"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"}, + "summary.jsonl": {"size": 51, "sha256": "1fcbf58d35ac56b969410719240d507f91ce4ccbcf9c68540f30c8adf225b439"} + }, + "checkpoint": { + "path": "checkpoints/steps_30000_pytorch_model.pt", + "size": 8456891339, + "sha256": "9646da2ae0b32589a75c8cc88fae96c93c5d269b69fd7a29200744936e01d96f" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 1073, + "vlm_tensors": 825, + "policy_tensors": 248, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4228247815, + "vlm_numel": 4073066496, + "policy_numel": 155181319, + "dtypes": { + "bfloat16": 1073 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [768, 256], + "action_model.model.transformer_blocks.0.attn1.to_k.weight": [768, 2048], + "action_model.model.transformer_blocks.1.attn1.to_k.weight": [768, 768], + "action_model.model.proj_out_2.weight": [1024, 768], + "action_model.action_decoder.layer2.weight": [7, 1024] + } + }, + "qwen25_pi": { + "model_type": "starvla", + "framework": "pi", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "oxe_bridge", + "directory": "qwen25-pi-bridge-rt1", + "repo_id": "StarVLA/Qwen-PI-Bridge-RT-1", + "revision": "26d0e079fbe3bc3fc62301f44f0025ef7c64ee22", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 3154, "sha256": "a7bdbde311bc910ee81e673a899199035d372641562824570c4dbaba4ea99ee2"}, + "dataset_statistics.json": {"size": 5990, "sha256": "83aa32682dd0b600e570936bfb63fd5d30b51d165e3d174a7ff3fc69d9fc276b"}, + "summary.jsonl": {"size": 169, "sha256": "9051086fe35e01366062e6e3eee43c54b081af37f12f7b132446b22bef129d70"} + }, + "checkpoint": { + "path": "checkpoints/steps_30000_pytorch_model.pt", + "size": 10103104403, + "sha256": "8a0e47858921924d5038f7c4393dee6682b83175a85546e35e357e8f74ce8343" + }, + "policy_prefixes": [ + "action_model." + ], + "expected": { + "total_tensors": 1073, + "vlm_tensors": 825, + "policy_tensors": 248, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 5051354119, + "vlm_numel": 4073066496, + "policy_numel": 978287623, + "total_nbytes": 10102708238, + "vlm_nbytes": 8146132992, + "policy_nbytes": 1956575246, + "dtypes": { + "bfloat16": 1073 + }, + "storage_alias_groups": 0 + }, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [2048, 256], + "action_model.model.transformer_blocks.0.attn1.to_k.weight": [2048, 2048], + "action_model.model.transformer_blocks.15.attn1.to_k.weight": [2048, 2048], + "action_model.model.proj_out_2.weight": [1024, 2048], + "action_model.action_decoder.layer2.weight": [7, 2048] + } + }, + "qwen25_fast": { + "model_type": "starvla", + "framework": "fast", + "backbone": "qwen2_5_vl", + "qwen_asset": "qwen2_5_vl_3b_instruct_action", + "default_unnorm_key": "bridge_dataset", + "directory": "qwen25-fast-bridge-rt1", + "repo_id": "StarVLA/Qwen-FAST-Bridge-RT-1", + "revision": "d9e2977d21755e78a0dd5f9a61586075a636d669", + "files": [ + "config.yaml", + "dataset_statistics.json", + "summary.jsonl" + ], + "file_hashes": { + "config.yaml": {"size": 2841, "sha256": "c0520794c8e5a15841b09fc1d9fb43216394674510d60e1f6d27dec149fa34f2"}, + "dataset_statistics.json": {"size": 6007, "sha256": "d2c4803c94d3b6eb1b8e8e100280e16a53b6058c8c2a7e747d27ecf9fcf9a4de"}, + "summary.jsonl": {"size": 34, "sha256": "2e927f0dd6524ec7cad6bb3023153142ee1c1f8015ba3712fcb0581dbf8c95e0"} + }, + "checkpoint": { + "path": "checkpoints/steps_10000_pytorch_model.pt", + "size": 8146439050, + "sha256": "f30e89a6b2a166fa3f48af42d5cffde07be44074b861abc7b57e1ccdb734e81e" + }, + "policy_prefixes": [], + "expected": null, + "required_shapes": { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048] + } + } + } +} diff --git a/tools/hf2gguf/starvla/convert.sh b/tools/hf2gguf/starvla/convert.sh new file mode 100755 index 0000000..7606055 --- /dev/null +++ b/tools/hf2gguf/starvla/convert.sh @@ -0,0 +1,160 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)" +PYTHON="${PYTHON:-${ROOT_DIR}/.venv/bin/python}" +CATALOG="${ROOT_DIR}/tools/hf2gguf/starvla/checkpoint_catalog.json" + +usage() { + cat >&2 <<'EOF' +Usage: tools/hf2gguf/starvla/convert.sh VARIANT [OUTPUT_DIR] [OPTIONS] + +Options: + --checkpoint PATH Convert a training checkpoint instead of the catalog release + --source-dir DIR Run directory containing config.yaml and dataset_statistics.json + --unnorm-key KEY Default normalization profile for a training checkpoint +EOF +} + +if [[ $# -eq 1 && ( "$1" == -h || "$1" == --help ) ]]; then + usage + exit 0 +fi +if [[ $# -lt 1 ]]; then + usage + exit 2 +fi + +VARIANT=$1 +shift +OUTPUT_DIR="${ROOT_DIR}/ckpts/starvla/gguf/${VARIANT}" +if [[ $# -gt 0 && "$1" != --* ]]; then + OUTPUT_DIR=$1 + shift +fi +CHECKPOINT_OVERRIDE="" +SOURCE_DIR_OVERRIDE="" +UNNORM_KEY="" +while [[ $# -gt 0 ]]; do + case "$1" in + --checkpoint|--source-dir|--unnorm-key) + [[ $# -ge 2 ]] || { echo "error: $1 requires a value" >&2; exit 2; } + case "$1" in + --checkpoint) CHECKPOINT_OVERRIDE=$2 ;; + --source-dir) SOURCE_DIR_OVERRIDE=$2 ;; + --unnorm-key) UNNORM_KEY=$2 ;; + esac + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "error: unknown option: $1" >&2 + usage + exit 2 + ;; + esac +done +if [[ -z "${CHECKPOINT_OVERRIDE}" && ( -n "${SOURCE_DIR_OVERRIDE}" || -n "${UNNORM_KEY}" ) ]]; then + echo "error: --source-dir and --unnorm-key require --checkpoint" >&2 + exit 2 +fi +[[ -x "${PYTHON}" ]] || { echo "error: missing Python: ${PYTHON}" >&2; exit 2; } +[[ ! -e "${OUTPUT_DIR}" ]] || { + echo "error: refusing to overwrite output directory: ${OUTPUT_DIR}" >&2 + exit 2 +} + +export STARVLA_CONFIG_PYTHON="${PYTHON}" +source "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_variant_config.sh" +load_starvla_variant "${VARIANT}" + +download_args=(--variant "${VARIANT}") +[[ "${FRAMEWORK}" == fast ]] && download_args+=(--include-fast-weights) +[[ -n "${CHECKPOINT_OVERRIDE}" ]] && download_args+=(--skip-checkpoint) +[[ "${STARVLA_LOCAL_FILES_ONLY:-0}" == 1 ]] && download_args+=(--local-files-only) +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/download_starvla.py" \ + --catalog "${CATALOG}" "${download_args[@]}" + +LLAMA_REV="$("${PYTHON}" -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["source_revisions"]["llama_cpp"])' \ + "${CATALOG}")" +export LLAMA_ROOT="${LLAMA_ROOT:-${ROOT_DIR}/ckpts/starvla/toolchains/llama.cpp-${LLAMA_REV}}" +if [[ ! -d "${LLAMA_ROOT}" ]]; then + mkdir -p -- "$(dirname -- "${LLAMA_ROOT}")" + git -C "${ROOT_DIR}/third_party/llama.cpp" worktree add \ + --detach "${LLAMA_ROOT}" "${LLAMA_REV}" +fi + +SOURCE_DIR="${ROOT_DIR}/ckpts/starvla/sources/${CHECKPOINT_DIRECTORY}/${CHECKPOINT_REVISION}" +CHECKPOINT="${SOURCE_DIR}/${CHECKPOINT_RELATIVE_PATH}" +BASE_ASSETS="${ROOT_DIR}/ckpts/starvla/sources/${QWEN_DIRECTORY}/${QWEN_REVISION}" +mkdir -p -- "${ROOT_DIR}/ckpts/starvla/work" +WORK_DIR="$(mktemp -d "${ROOT_DIR}/ckpts/starvla/work/.${VARIANT}.XXXXXX")" +cleanup() { rm -rf -- "${WORK_DIR}"; } +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +if [[ -n "${CHECKPOINT_OVERRIDE}" ]]; then + CHECKPOINT="$(realpath -- "${CHECKPOINT_OVERRIDE}")" + if [[ -n "${SOURCE_DIR_OVERRIDE}" ]]; then + SOURCE_DIR="$(realpath -- "${SOURCE_DIR_OVERRIDE}")" + else + checkpoint_dir="$(dirname -- "${CHECKPOINT}")" + for candidate in "${checkpoint_dir}" "$(dirname -- "${checkpoint_dir}")"; do + if [[ -f "${candidate}/config.yaml" && -f "${candidate}/dataset_statistics.json" ]]; then + SOURCE_DIR="${candidate}" + break + fi + done + fi + [[ -f "${SOURCE_DIR}/config.yaml" && -f "${SOURCE_DIR}/dataset_statistics.json" ]] || { + echo "error: cannot find config.yaml and dataset_statistics.json; pass --source-dir" >&2 + exit 2 + } + LOCAL_CATALOG="${WORK_DIR}/checkpoint_catalog.json" + "${PYTHON}" - "${CATALOG}" "${LOCAL_CATALOG}" "${VARIANT}" \ + "${CHECKPOINT}" "${SOURCE_DIR}" "${UNNORM_KEY}" <<'PY' +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(sys.argv[1]).parent)) +from starvla_checkpoint import atomic_write_json, load_catalog, local_checkpoint_catalog + +catalog = local_checkpoint_catalog( + load_catalog(Path(sys.argv[1])), + sys.argv[3], + Path(sys.argv[4]), + Path(sys.argv[5]), + sys.argv[6] or None, +) +atomic_write_json(Path(sys.argv[2]), catalog) +PY + CATALOG="${LOCAL_CATALOG}" + export STARVLA_CATALOG="${CATALOG}" + load_starvla_variant "${VARIANT}" +fi + +export PYTHON VARIANT SOURCE_DIR CHECKPOINT BASE_ASSETS WORK_DIR OUTPUT_DIR LLAMA_ROOT +if [[ "${FRAMEWORK}" != fast ]]; then + bash "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_all.sh" +else + CODEC_REV="$("${PYTHON}" -c \ + 'import json,sys; print(json.load(open(sys.argv[1]))["shared_assets"]["fast_codec"]["revision"])' \ + "${CATALOG}")" + "${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py" \ + --checkpoint "${CHECKPOINT}" \ + --source-dir "${SOURCE_DIR}" \ + --qwen-assets "${BASE_ASSETS}" \ + --fast-codec "${ROOT_DIR}/ckpts/starvla/sources/fast-codec/${CODEC_REV}" \ + --staging-dir "${WORK_DIR}/staging" \ + --output-dir "${OUTPUT_DIR}" \ + --catalog "${CATALOG}" \ + --llama-root "${LLAMA_ROOT}" \ + --python "${PYTHON}" +fi + +echo "StarVLA ${VARIANT} bundle: ${OUTPUT_DIR}" diff --git a/tools/hf2gguf/starvla/convert_starvla_all.sh b/tools/hf2gguf/starvla/convert_starvla_all.sh new file mode 100755 index 0000000..a0fa6c5 --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_all.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +PYTHON="${PYTHON:-${ROOT_DIR}/.venv/bin/python}" +CATALOG="${STARVLA_CATALOG:-${ROOT_DIR}/tools/hf2gguf/starvla/checkpoint_catalog.json}" +LLAMA_ROOT="${LLAMA_ROOT:?set LLAMA_ROOT to an absolute clean checkout of the catalog-pinned llama.cpp revision}" +source "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_variant_config.sh" +VARIANT="${VARIANT:?set VARIANT to oft, groot, pi_v3, qwen25_oft, qwen25_groot, or qwen25_pi}" +load_starvla_variant "${VARIANT}" + +if [[ "${FRAMEWORK}" == "fast" ]]; then + echo "error: FAST uses tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py" >&2 + exit 1 +fi + +CHECKPOINT="${CHECKPOINT:?set CHECKPOINT to the pinned ${VARIANT} .pt file}" +SOURCE_DIR="${SOURCE_DIR:?set SOURCE_DIR to the pinned ${VARIANT} source directory}" +BASE_ASSETS="${BASE_ASSETS:?set BASE_ASSETS to the pinned Qwen-VL asset directory}" +WORK_DIR="${WORK_DIR:-${ROOT_DIR}/ckpts/starvla/work/${VARIANT}}" +OUTPUT_DIR="${OUTPUT_DIR:-${ROOT_DIR}/ckpts/starvla/gguf/${VARIANT}}" +MAX_SHARD_SIZE="${MAX_SHARD_SIZE:-2G}" +TEXT_DTYPE="${TEXT_DTYPE:-bf16}" +MMPROJ_DTYPE="${MMPROJ_DTYPE:-bf16}" +POLICY_DTYPE="${POLICY_DTYPE:-fp32}" +TEXT_FILENAME="${TEXT_FILENAME:-qwen-${ARTIFACT_STEM}-${TEXT_DTYPE}.gguf}" +MMPROJ_FILENAME="${MMPROJ_FILENAME:-mmproj-${ARTIFACT_STEM}-${MMPROJ_DTYPE}.gguf}" +POLICY_FILENAME="${POLICY_FILENAME:-starvla-${ARTIFACT_STEM}-policy-${POLICY_DTYPE}.gguf}" +MANIFEST_FILENAME="conversion_manifest.json" + +if [[ -e "${WORK_DIR}" ]] && + [[ -n "$(find "${WORK_DIR}" -mindepth 1 -maxdepth 1 -print -quit)" ]]; then + echo "error: WORK_DIR must be empty: ${WORK_DIR}" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" + +for filename in \ + "${TEXT_FILENAME}" \ + "${MMPROJ_FILENAME}" \ + "${POLICY_FILENAME}" \ + "${MANIFEST_FILENAME}"; do + destination="${OUTPUT_DIR}/${filename}" + if [[ -e "${destination}" || -L "${destination}" ]]; then + echo "error: refusing to overwrite existing output: ${destination}" >&2 + exit 1 + fi +done + +RUN_OUTPUT_DIR="" +declare -a PUBLISHED_FILES=() +SUCCESS=0 + +cleanup() { + status=$? + trap - EXIT + set +e + if [[ -n "${RUN_OUTPUT_DIR}" ]]; then + rm -rf -- "${RUN_OUTPUT_DIR}" + fi + if [[ "${SUCCESS}" != 1 ]]; then + for published in "${PUBLISHED_FILES[@]}"; do + rm -f -- "${published}" + done + rm -rf -- "${WORK_DIR}" + fi + exit "${status}" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +mkdir -p "${WORK_DIR}" +RUN_OUTPUT_DIR="$(mktemp -d "${OUTPUT_DIR}/.starvla-${ARTIFACT_STEM}.tmp.XXXXXX")" + +publish_file() { + source=$1 + destination=$2 + if [[ ! -f "${source}" || ! -s "${source}" ]]; then + echo "error: transaction output is missing or empty: ${source}" >&2 + return 1 + fi + if ! ln -- "${source}" "${destination}"; then + echo "error: refusing to overwrite existing output: ${destination}" >&2 + return 1 + fi + PUBLISHED_FILES+=("${destination}") + rm -f -- "${source}" +} + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/starvla_surgery.py" \ + "${CHECKPOINT}" \ + --variant "${VARIANT}" \ + --catalog "${CATALOG}" \ + --source-dir "${SOURCE_DIR}" \ + --base-assets "${BASE_ASSETS}" \ + --output-dir "${WORK_DIR}/staging" \ + --max-shard-size "${MAX_SHARD_SIZE}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --output-dir "${RUN_OUTPUT_DIR}" \ + --llama-root "${LLAMA_ROOT}" \ + --text-filename "${TEXT_FILENAME}" \ + --mmproj-filename "${MMPROJ_FILENAME}" \ + --text-dtype "${TEXT_DTYPE}" \ + --mmproj-dtype "${MMPROJ_DTYPE}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py" \ + --variant "${VARIANT}" \ + --policy-dir "${WORK_DIR}/staging/policy" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --output "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" \ + --dtype "${POLICY_DTYPE}" \ + --text-filename "${TEXT_FILENAME}" \ + --mmproj-filename "${MMPROJ_FILENAME}" + +"${PYTHON}" "${ROOT_DIR}/tools/hf2gguf/starvla/validate_starvla_bundle.py" \ + --variant "${VARIANT}" \ + --text "${RUN_OUTPUT_DIR}/${TEXT_FILENAME}" \ + --mmproj "${RUN_OUTPUT_DIR}/${MMPROJ_FILENAME}" \ + --policy "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" \ + --hf-dir "${WORK_DIR}/staging/hf" \ + --policy-dir "${WORK_DIR}/staging/policy" \ + --surgery-manifest "${WORK_DIR}/staging/surgery_manifest.json" \ + --catalog "${CATALOG}" \ + --text-dtype "${TEXT_DTYPE}" \ + --mmproj-dtype "${MMPROJ_DTYPE}" \ + --policy-dtype "${POLICY_DTYPE}" \ + --output "${RUN_OUTPUT_DIR}/${MANIFEST_FILENAME}" + +publish_file "${RUN_OUTPUT_DIR}/${TEXT_FILENAME}" "${OUTPUT_DIR}/${TEXT_FILENAME}" +publish_file "${RUN_OUTPUT_DIR}/${MMPROJ_FILENAME}" "${OUTPUT_DIR}/${MMPROJ_FILENAME}" +publish_file "${RUN_OUTPUT_DIR}/${POLICY_FILENAME}" "${OUTPUT_DIR}/${POLICY_FILENAME}" +# The manifest is the bundle commit marker and is intentionally published last. +publish_file "${RUN_OUTPUT_DIR}/${MANIFEST_FILENAME}" "${OUTPUT_DIR}/${MANIFEST_FILENAME}" +SUCCESS=1 + +echo "StarVLA ${VARIANT} bundle written to ${OUTPUT_DIR}" diff --git a/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py new file mode 100755 index 0000000..21dd35a --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_policy_to_gguf.py @@ -0,0 +1,1819 @@ +#!/usr/bin/env python3 +"""Convert a StarVLA policy staging directory to a policy GGUF.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +from pathlib import Path +from typing import Any, Mapping + +import numpy as np + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_POLICY_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + create_output_temporary, + default_mmproj_filename, + default_text_filename, + get_variant, + load_catalog, + resolve_effective_config, + sha256_file, + validate_surgery_manifest, + verify_staged_assets, +) + + +OFT_TENSOR_MAP = { + "action_model.model.layer_norm1.weight": "starvla.policy.oft.input_norm.weight", + "action_model.model.layer_norm1.bias": "starvla.policy.oft.input_norm.bias", + "action_model.model.fc1.weight": "starvla.policy.oft.input_proj.weight", + "action_model.model.fc1.bias": "starvla.policy.oft.input_proj.bias", + "action_model.model.mlp_resnet_blocks.0.ffn.0.weight": "starvla.policy.oft.block.0.norm.weight", + "action_model.model.mlp_resnet_blocks.0.ffn.0.bias": "starvla.policy.oft.block.0.norm.bias", + "action_model.model.mlp_resnet_blocks.0.ffn.1.weight": "starvla.policy.oft.block.0.linear.weight", + "action_model.model.mlp_resnet_blocks.0.ffn.1.bias": "starvla.policy.oft.block.0.linear.bias", + "action_model.model.mlp_resnet_blocks.1.ffn.0.weight": "starvla.policy.oft.block.1.norm.weight", + "action_model.model.mlp_resnet_blocks.1.ffn.0.bias": "starvla.policy.oft.block.1.norm.bias", + "action_model.model.mlp_resnet_blocks.1.ffn.1.weight": "starvla.policy.oft.block.1.linear.weight", + "action_model.model.mlp_resnet_blocks.1.ffn.1.bias": "starvla.policy.oft.block.1.linear.bias", + "action_model.model.layer_norm2.weight": "starvla.policy.oft.output_norm.weight", + "action_model.model.layer_norm2.bias": "starvla.policy.oft.output_norm.bias", + "action_model.model.fc2.weight": "starvla.policy.oft.output_proj.weight", + "action_model.model.fc2.bias": "starvla.policy.oft.output_proj.bias", +} + + +DIT_BLOCK_SUFFIXES = { + "norm1.linear.weight": "ada_norm.weight", + "norm1.linear.bias": "ada_norm.bias", + "attn1.to_q.weight": "attention.query.weight", + "attn1.to_q.bias": "attention.query.bias", + "attn1.to_k.weight": "attention.key.weight", + "attn1.to_k.bias": "attention.key.bias", + "attn1.to_v.weight": "attention.value.weight", + "attn1.to_v.bias": "attention.value.bias", + "attn1.to_out.0.weight": "attention.output.weight", + "attn1.to_out.0.bias": "attention.output.bias", + "ff.net.0.proj.weight": "feed_forward.input.weight", + "ff.net.0.proj.bias": "feed_forward.input.bias", + "ff.net.2.weight": "feed_forward.output.weight", + "ff.net.2.bias": "feed_forward.output.bias", +} + + +def _build_flow_tensor_map(framework: str, block_count: int) -> dict[str, str]: + destination = f"starvla.policy.{framework}" + tensor_map = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": + f"{destination}.timestep.input.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": + f"{destination}.timestep.input.bias", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": + f"{destination}.timestep.output.weight", + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": + f"{destination}.timestep.output.bias", + "action_model.action_encoder.layer1.weight": f"{destination}.action.input.weight", + "action_model.action_encoder.layer1.bias": f"{destination}.action.input.bias", + "action_model.action_encoder.layer2.weight": f"{destination}.action.time_mix.weight", + "action_model.action_encoder.layer2.bias": f"{destination}.action.time_mix.bias", + "action_model.action_encoder.layer3.weight": f"{destination}.action.output.weight", + "action_model.action_encoder.layer3.bias": f"{destination}.action.output.bias", + "action_model.action_decoder.layer1.weight": f"{destination}.velocity.input.weight", + "action_model.action_decoder.layer1.bias": f"{destination}.velocity.input.bias", + "action_model.action_decoder.layer2.weight": f"{destination}.velocity.output.weight", + "action_model.action_decoder.layer2.bias": f"{destination}.velocity.output.bias", + "action_model.future_tokens.weight": f"{destination}.future_tokens.weight", + "action_model.position_embedding.weight": f"{destination}.action_position.weight", + } + for block in range(block_count): + source = f"action_model.model.transformer_blocks.{block}" + target = f"{destination}.block.{block}" + for source_suffix, destination_suffix in DIT_BLOCK_SUFFIXES.items(): + tensor_map[f"{source}.{source_suffix}"] = f"{target}.{destination_suffix}" + return tensor_map + + +def build_groot_tensor_map(block_count: int = 16) -> dict[str, str]: + tensor_map = _build_flow_tensor_map("groot", block_count) + tensor_map.update( + { + "action_model.model.proj_out_1.weight": "starvla.policy.groot.output.modulation.weight", + "action_model.model.proj_out_1.bias": "starvla.policy.groot.output.modulation.bias", + "action_model.model.proj_out_2.weight": "starvla.policy.groot.output.projection.weight", + "action_model.model.proj_out_2.bias": "starvla.policy.groot.output.projection.bias", + } + ) + return tensor_map + + +def build_pi_tensor_map(block_count: int = 16) -> dict[str, str]: + tensor_map = _build_flow_tensor_map("pi", block_count) + tensor_map.update( + { + "action_model.state_encoder.layer1.weight": "starvla.policy.pi.state.input.weight", + "action_model.state_encoder.layer1.bias": "starvla.policy.pi.state.input.bias", + "action_model.state_encoder.layer2.weight": "starvla.policy.pi.state.output.weight", + "action_model.state_encoder.layer2.bias": "starvla.policy.pi.state.output.bias", + } + ) + return tensor_map + + +def build_pi_v3_tensor_map( + block_count: int = 36, + projector_count: int = 36, +) -> dict[str, str]: + tensor_map = _build_flow_tensor_map("pi_v3", block_count) + for projector in range(projector_count): + source = f"project_layers.{projector}" + target = f"starvla.policy.pi_v3.projector.{projector}" + tensor_map.update( + { + f"{source}.0.weight": f"{target}.norm.weight", + f"{source}.0.bias": f"{target}.norm.bias", + f"{source}.1.weight": f"{target}.projection.weight", + f"{source}.1.bias": f"{target}.projection.bias", + } + ) + return tensor_map + + +GROOT_BLOCK_COUNT = 16 +GROOT_TENSOR_MAP = build_groot_tensor_map(GROOT_BLOCK_COUNT) +GROOT_UNUSED_SOURCE_TENSORS = { + "action_model.state_encoder.layer1.weight", + "action_model.state_encoder.layer1.bias", + "action_model.state_encoder.layer2.weight", + "action_model.state_encoder.layer2.bias", +} +GROOT_SOURCE_TENSOR_COUNT = 248 +GROOT_POLICY_TENSOR_COUNT = 244 +GROOT_QWEN3_POLICY_NUMEL = 161_472_775 +GROOT_QWEN25_POLICY_NUMEL = 155_181_319 +GROOT_DIT_NORM_EPS = 1e-5 +GROOT_OUTPUT_NORM_EPS = 1e-6 +GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE = { + backbone: { + "qwen_hidden_dim": qwen_hidden_dim, + "dit_width": 768, + "timestep_dim": 256, + "feed_forward_dim": 3072, + "output_dim": 1024, + "mlp_hidden_dim": 1024, + "state_dim": 7, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": GROOT_BLOCK_COUNT, + "tensor_count": GROOT_SOURCE_TENSOR_COUNT, + "numel": numel, + } + for backbone, qwen_hidden_dim, numel in ( + ("qwen3_vl", 2560, GROOT_QWEN3_POLICY_NUMEL), + ("qwen2_5_vl", 2048, GROOT_QWEN25_POLICY_NUMEL), + ) +} + +PI_BLOCK_COUNT = 16 +PI_TENSOR_MAP = build_pi_tensor_map(PI_BLOCK_COUNT) +PI_UNUSED_SOURCE_TENSORS = { + "action_model.model.proj_out_1.weight", + "action_model.model.proj_out_1.bias", + "action_model.model.proj_out_2.weight", + "action_model.model.proj_out_2.bias", +} +PI_POLICY_TENSOR_COUNT = 244 +PI_POLICY_NUMEL = 967_796_743 +PI_DIT_NORM_EPS = 1e-5 +PI_SUPPORTED_DIMENSIONS = { + "qwen_hidden_dim": 2048, + "dit_width": 2048, + "timestep_dim": 256, + "feed_forward_dim": 8192, + "mlp_hidden_dim": 2048, + "state_dim": 7, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": PI_BLOCK_COUNT, + "tensor_count": PI_POLICY_TENSOR_COUNT, + "numel": PI_POLICY_NUMEL, +} + +PI_V3_BLOCK_COUNT = 36 +PI_V3_PROJECTOR_COUNT = 36 +PI_V3_TENSOR_MAP = build_pi_v3_tensor_map(PI_V3_BLOCK_COUNT, PI_V3_PROJECTOR_COUNT) +PI_V3_POLICY_TENSOR_COUNT = len(PI_V3_TENSOR_MAP) +PI_V3_DIT_NORM_EPS = 1e-5 +PI_V3_PROJECTOR_NORM_EPS = 1e-5 +PI_V3_SUPPORTED_DIMENSIONS = { + "qwen_hidden_dim": 2560, + "dit_width": 1024, + "timestep_dim": 256, + "feed_forward_dim": 4096, + "mlp_hidden_dim": 1024, + "action_dim": 7, + "future_token_count": 32, + "max_sequence_length": 1024, + "block_count": PI_V3_BLOCK_COUNT, + "projector_count": PI_V3_PROJECTOR_COUNT, + "tensor_count": PI_V3_POLICY_TENSOR_COUNT, +} + +ACTION_NAMES = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] +OFT_ACTION_TOKEN = "🔍" +OFT_ACTION_TOKEN_ID = 146663 +OFT_LAYER_NORM_EPS = 1e-5 +QWEN3VL_PROCESSOR_MIN_PIXELS = 65_536 +QWEN3VL_PROCESSOR_MAX_PIXELS = 16_777_216 +QWEN3VL_IMAGE_PATCH_SIZE = 16 +QWEN3VL_TEMPORAL_PATCH_SIZE = 2 +QWEN3VL_SPATIAL_MERGE_SIZE = 2 +QWEN3VL_MIN_IMAGE_TOKENS = 64 +QWEN3VL_MAX_IMAGE_TOKENS = 16_384 +QWEN3VL_IMAGE_MEAN = [0.5, 0.5, 0.5] +QWEN3VL_IMAGE_STD = [0.5, 0.5, 0.5] +QWEN25VL_PROCESSOR_MIN_PIXELS = 3_136 +QWEN25VL_PROCESSOR_MAX_PIXELS = 12_845_056 +QWEN25VL_IMAGE_PATCH_SIZE = 14 +QWEN25VL_TEMPORAL_PATCH_SIZE = 2 +QWEN25VL_SPATIAL_MERGE_SIZE = 2 +QWEN25VL_MIN_IMAGE_TOKENS = 4 +QWEN25VL_MAX_IMAGE_TOKENS = 16_384 +QWEN25VL_IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] +QWEN25VL_IMAGE_STD = [0.26862954, 0.26130258, 0.27577711] +# These defaults are executable behavior in the pinned Transformers 4.57 fast processor, +# including antialias=True on its torchvision resize call. +QWEN3VL_DYNAMIC_IMAGE_METADATA = { + "starvla.image.count": 1, + "starvla.image.names": ["image_0"], + "starvla.image.preprocessing_mode": "qwen3vl_smart_resize", + "starvla.image.framework_inference_pre_resize": False, + "starvla.image.framework_inference_pre_resize_config_key": ( + "datasets.vla_data.obs_image_size" + ), + "starvla.image.processor_min_pixels": QWEN3VL_PROCESSOR_MIN_PIXELS, + "starvla.image.processor_max_pixels": QWEN3VL_PROCESSOR_MAX_PIXELS, + "starvla.image.processor_class": "Qwen2VLImageProcessorFast", + "starvla.image.processor_reference_transformers_version": "4.57.0", + "starvla.image.processor_do_convert_rgb": True, + "starvla.image.processor_do_resize": True, + "starvla.image.processor_resize_resample": "bicubic", + "starvla.image.processor_resize_antialias": True, + "starvla.image.processor_do_rescale": True, + "starvla.image.processor_rescale_factor": 1.0 / 255.0, + "starvla.image.processor_do_normalize": True, + "starvla.image.processor_image_mean": QWEN3VL_IMAGE_MEAN, + "starvla.image.processor_image_std": QWEN3VL_IMAGE_STD, + "starvla.image.patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "starvla.image.temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "starvla.image.spatial_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "starvla.image.token_count_mode": "dynamic_grid_thw_after_spatial_merge", + "starvla.image.min_token_count": QWEN3VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN3VL_MAX_IMAGE_TOKENS, +} + + +def _load_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def _write_gguf_arrays_no_overwrite( + output: Path, + metadata: dict[str, Any], + arrays: Any, + writer: Any, +) -> None: + """Write beside the destination, then publish without replacing an existing file.""" + output.parent.mkdir(parents=True, exist_ok=True) + if output.exists() or output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {output}") + + descriptor, temporary = create_output_temporary(output) + os.close(descriptor) + try: + writer(temporary, metadata, arrays) + if not temporary.is_file() or temporary.stat().st_size == 0: + raise StarVLAError(f"GGUF writer did not create a non-empty output: {temporary}") + try: + os.link(temporary, output) + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output: {output}") from exc + temporary.unlink() + finally: + temporary.unlink(missing_ok=True) + + +def _load_yaml(path: Path) -> dict[str, Any]: + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to load a StarVLA policy config") from exc + try: + value = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load YAML {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a YAML object in {path}") + return value + + +def load_policy_tensors(policy_dir: Path) -> dict[str, Any]: + try: + from safetensors import safe_open + except ImportError as exc: + raise StarVLAError("safetensors is required to convert a StarVLA policy") from exc + + index_path = policy_dir / "policy.safetensors.index.json" + index = _load_json(index_path) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise StarVLAError(f"invalid or empty safetensors weight_map in {index_path}") + + tensors = {} + by_shard: dict[str, list[str]] = {} + for name, shard in weight_map.items(): + by_shard.setdefault(str(shard), []).append(str(name)) + for shard, names in sorted(by_shard.items()): + shard_path = policy_dir / shard + if not shard_path.is_file(): + raise StarVLAError(f"missing policy shard: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as handle: + if set(handle.keys()) != set(names): + raise StarVLAError(f"policy shard/index key mismatch: {shard_path}") + for name in sorted(names): + tensors[name] = handle.get_tensor(name) + if set(tensors) != set(weight_map): + raise StarVLAError("loaded policy tensor set does not match the index") + return tensors + + +def _tensor_shape(tensors: Mapping[str, Any], name: str) -> list[int]: + return [int(dimension) for dimension in tensors[name].shape] + + +def _matrix_shape(tensors: Mapping[str, Any], name: str) -> list[int]: + shape = _tensor_shape(tensors, name) + if len(shape) != 2: + raise StarVLAError(f"invalid matrix shape for {name}: {shape}") + return shape + + +def _validate_tensor_shapes( + tensors: Mapping[str, Any], expected: Mapping[str, list[int]], *, label: str +) -> None: + mismatches = [ + f"{name}: expected {shape}, got {_tensor_shape(tensors, name)}" + for name, shape in expected.items() + if _tensor_shape(tensors, name) != shape + ] + if mismatches: + raise StarVLAError(f"invalid {label} tensor shapes: " + "; ".join(mismatches)) + + +def _dit_block_shapes( + prefix: str, width: int, attention_dim: int, feed_forward_dim: int +) -> dict[str, list[int]]: + return { + f"{prefix}.norm1.linear.weight": [2 * width, width], + f"{prefix}.norm1.linear.bias": [2 * width], + f"{prefix}.attn1.to_q.weight": [width, width], + f"{prefix}.attn1.to_q.bias": [width], + f"{prefix}.attn1.to_k.weight": [width, attention_dim], + f"{prefix}.attn1.to_k.bias": [width], + f"{prefix}.attn1.to_v.weight": [width, attention_dim], + f"{prefix}.attn1.to_v.bias": [width], + f"{prefix}.attn1.to_out.0.weight": [width, width], + f"{prefix}.attn1.to_out.0.bias": [width], + f"{prefix}.ff.net.0.proj.weight": [feed_forward_dim, width], + f"{prefix}.ff.net.0.proj.bias": [feed_forward_dim], + f"{prefix}.ff.net.2.weight": [width, feed_forward_dim], + f"{prefix}.ff.net.2.bias": [width], + } + + +def validate_oft_tensors(tensors: dict[str, Any]) -> dict[str, int]: + actual = set(tensors) + expected = set(OFT_TENSOR_MAP) + if actual != expected: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected) + raise StarVLAError(f"OFT policy tensor mismatch; missing={missing}, unexpected={unexpected}") + + input_dim = _tensor_shape(tensors, "action_model.model.layer_norm1.weight")[0] + input_projection = _tensor_shape(tensors, "action_model.model.fc1.weight") + if len(input_projection) != 2 or input_projection[1] != input_dim: + raise StarVLAError(f"invalid OFT input projection shape: {input_projection}") + hidden_dim = input_projection[0] + output_projection = _tensor_shape(tensors, "action_model.model.fc2.weight") + if len(output_projection) != 2 or output_projection[1] != hidden_dim: + raise StarVLAError(f"invalid OFT output projection shape: {output_projection}") + action_dim = output_projection[0] + + expected_shapes = { + "action_model.model.layer_norm1.weight": [input_dim], + "action_model.model.layer_norm1.bias": [input_dim], + "action_model.model.fc1.weight": [hidden_dim, input_dim], + "action_model.model.fc1.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.0.weight": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.0.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.1.weight": [hidden_dim, hidden_dim], + "action_model.model.mlp_resnet_blocks.0.ffn.1.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.0.weight": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.0.bias": [hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.1.weight": [hidden_dim, hidden_dim], + "action_model.model.mlp_resnet_blocks.1.ffn.1.bias": [hidden_dim], + "action_model.model.layer_norm2.weight": [hidden_dim], + "action_model.model.layer_norm2.bias": [hidden_dim], + "action_model.model.fc2.weight": [action_dim, hidden_dim], + "action_model.model.fc2.bias": [action_dim], + } + _validate_tensor_shapes(tensors, expected_shapes, label="OFT") + return {"input_dim": input_dim, "hidden_dim": hidden_dim, "action_dim": action_dim} + + +def validate_groot_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate GR00T policy tensors and infer their dimensions.""" + actual = set(tensors) + expected = set(GROOT_TENSOR_MAP) + if not expected.issubset(actual) or actual - expected != GROOT_UNUSED_SOURCE_TENSORS: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected - GROOT_UNUSED_SOURCE_TENSORS) + raise StarVLAError(f"GR00T policy tensor mismatch; missing={missing}, unexpected={unexpected}") + + timestep_input = _matrix_shape(tensors, + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + cross_attention_dim = _matrix_shape(tensors, + "action_model.model.transformer_blocks.0.attn1.to_k.weight" + )[1] + feed_forward_dim = _matrix_shape(tensors, + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + output_dim = _matrix_shape(tensors, "action_model.model.proj_out_2.weight")[0] + mlp_hidden_dim = _matrix_shape(tensors, "action_model.action_decoder.layer1.weight")[0] + state_dim = _matrix_shape(tensors, "action_model.state_encoder.layer1.weight")[1] + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.model.proj_out_1.weight": [2 * dit_width, dit_width], + "action_model.model.proj_out_1.bias": [2 * dit_width], + "action_model.model.proj_out_2.weight": [output_dim, dit_width], + "action_model.model.proj_out_2.bias": [output_dim], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, output_dim], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(GROOT_BLOCK_COUNT): + attention_dim = cross_attention_dim if block % 2 == 0 else dit_width + expected_shapes.update( + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + attention_dim, + feed_forward_dim, + ) + ) + _validate_tensor_shapes(tensors, expected_shapes, label="GR00T") + + numel = sum(int(tensor.numel()) for tensor in tensors.values()) + return { + "qwen_hidden_dim": cross_attention_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "output_dim": output_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "state_dim": state_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": GROOT_BLOCK_COUNT, + "tensor_count": len(tensors), + "numel": numel, + } + + +def validate_pi_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate the tensors used by the legacy Qwen-PI inference graph.""" + actual = set(tensors) + expected = set(PI_TENSOR_MAP) + if not expected.issubset(actual) or actual - expected != PI_UNUSED_SOURCE_TENSORS: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected - PI_UNUSED_SOURCE_TENSORS) + raise StarVLAError( + f"legacy PI policy tensor mismatch; missing={missing}, unexpected={unexpected}" + ) + + timestep_input = _matrix_shape(tensors, + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + cross_attention_dim = _matrix_shape(tensors, + "action_model.model.transformer_blocks.0.attn1.to_k.weight" + )[1] + feed_forward_dim = _matrix_shape(tensors, + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + mlp_hidden_dim, state_dim = _matrix_shape(tensors, "action_model.state_encoder.layer1.weight") + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.state_encoder.layer1.weight": [mlp_hidden_dim, state_dim], + "action_model.state_encoder.layer1.bias": [mlp_hidden_dim], + "action_model.state_encoder.layer2.weight": [dit_width, mlp_hidden_dim], + "action_model.state_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, dit_width], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(PI_BLOCK_COUNT): + expected_shapes.update( + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + cross_attention_dim, + feed_forward_dim, + ) + ) + _validate_tensor_shapes(tensors, expected_shapes, label="legacy PI") + + return { + "qwen_hidden_dim": cross_attention_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "state_dim": state_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": PI_BLOCK_COUNT, + "tensor_count": len(expected), + "numel": sum(int(tensors[name].numel()) for name in expected), + } + + +def validate_pi_v3_tensors(tensors: dict[str, Any]) -> dict[str, int]: + """Validate PI_v3 policy tensors and infer their dimensions.""" + actual = set(tensors) + expected = set(PI_V3_TENSOR_MAP) + missing = sorted(expected - actual) + if missing: + raise StarVLAError(f"PI-v3 policy is missing runtime tensors: {missing}") + + timestep_input = _matrix_shape(tensors, + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight" + ) + dit_width, timestep_dim = timestep_input + feed_forward_dim = _matrix_shape(tensors, + "action_model.model.transformer_blocks.0.ff.net.0.proj.weight" + )[0] + mlp_hidden_dim = _matrix_shape(tensors, "action_model.action_decoder.layer1.weight")[0] + action_dim = _matrix_shape(tensors, "action_model.action_encoder.layer1.weight")[1] + future_token_count = _matrix_shape(tensors, "action_model.future_tokens.weight")[0] + max_sequence_length = _matrix_shape(tensors, "action_model.position_embedding.weight")[0] + qwen_hidden_dim = _tensor_shape(tensors, "project_layers.0.0.weight")[0] + projector_output_dim = _matrix_shape(tensors, "project_layers.0.1.weight")[0] + if projector_output_dim != dit_width: + raise StarVLAError( + "invalid PI_v3 projector/DiT width contract: " + f"projector={projector_output_dim}, DiT={dit_width}" + ) + + expected_shapes = { + "action_model.model.timestep_encoder.timestep_embedder.linear_1.weight": [ + dit_width, + timestep_dim, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_1.bias": [dit_width], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.weight": [ + dit_width, + dit_width, + ], + "action_model.model.timestep_encoder.timestep_embedder.linear_2.bias": [dit_width], + "action_model.action_encoder.layer1.weight": [dit_width, action_dim], + "action_model.action_encoder.layer1.bias": [dit_width], + "action_model.action_encoder.layer2.weight": [dit_width, 2 * dit_width], + "action_model.action_encoder.layer2.bias": [dit_width], + "action_model.action_encoder.layer3.weight": [dit_width, dit_width], + "action_model.action_encoder.layer3.bias": [dit_width], + "action_model.action_decoder.layer1.weight": [mlp_hidden_dim, dit_width], + "action_model.action_decoder.layer1.bias": [mlp_hidden_dim], + "action_model.action_decoder.layer2.weight": [action_dim, mlp_hidden_dim], + "action_model.action_decoder.layer2.bias": [action_dim], + "action_model.future_tokens.weight": [future_token_count, dit_width], + "action_model.position_embedding.weight": [max_sequence_length, dit_width], + } + for block in range(PI_V3_BLOCK_COUNT): + expected_shapes.update( + _dit_block_shapes( + f"action_model.model.transformer_blocks.{block}", + dit_width, + dit_width, + feed_forward_dim, + ) + ) + for projector in range(PI_V3_PROJECTOR_COUNT): + prefix = f"project_layers.{projector}" + expected_shapes.update( + { + f"{prefix}.0.weight": [qwen_hidden_dim], + f"{prefix}.0.bias": [qwen_hidden_dim], + f"{prefix}.1.weight": [dit_width, qwen_hidden_dim], + f"{prefix}.1.bias": [dit_width], + } + ) + _validate_tensor_shapes(tensors, expected_shapes, label="PI_v3") + + return { + "qwen_hidden_dim": qwen_hidden_dim, + "dit_width": dit_width, + "timestep_dim": timestep_dim, + "feed_forward_dim": feed_forward_dim, + "mlp_hidden_dim": mlp_hidden_dim, + "action_dim": action_dim, + "future_token_count": future_token_count, + "max_sequence_length": max_sequence_length, + "block_count": PI_V3_BLOCK_COUNT, + "projector_count": PI_V3_PROJECTOR_COUNT, + "tensor_count": len(PI_V3_TENSOR_MAP), + } + + +def load_variant_config( + policy_dir: Path, + surgery_manifest: dict[str, Any], + variant_name: str, + backbone: str | None = None, +) -> dict[str, Any]: + catalog_variant = str(surgery_manifest.get("variant", variant_name)) + effective_backbone = backbone or surgery_manifest.get("backbone") + if not isinstance(effective_backbone, str): + raise StarVLAError("surgery manifest does not identify the Qwen backbone") + effective = resolve_effective_config( + policy_dir, + catalog_variant, + { + "framework": surgery_manifest.get("framework", variant_name), + "backbone": effective_backbone, + }, + ) + effective_path = policy_dir / "effective_config.json" + effective_record = surgery_manifest.get("effective_config", {}) + if not effective_path.is_file(): + raise StarVLAError(f"missing surgery effective config: {effective_path}") + if ( + effective_record.get("path") != effective_path.name + or effective_record.get("size") != effective_path.stat().st_size + or effective_record.get("sha256") != sha256_file(effective_path) + ): + raise StarVLAError( + f"effective {variant_name.upper()} config does not match its canonical source/manifest" + ) + stored_effective = _load_json(effective_path) + if stored_effective != effective: + raise StarVLAError( + f"effective {variant_name.upper()} config does not match its canonical source/manifest" + ) + return effective + + +def load_oft_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "oft", backbone) + + +def load_groot_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "groot", backbone) + + +def load_pi_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi", backbone) + + +def load_pi_v3_config( + policy_dir: Path, surgery_manifest: dict[str, Any], backbone: str +) -> dict[str, Any]: + return load_variant_config(policy_dir, surgery_manifest, "pi_v3", backbone) + + +def resolve_action_token_id(hf_dir: Path) -> int: + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise StarVLAError("transformers is required to verify the OFT action token") from exc + try: + tokenizer = AutoTokenizer.from_pretrained(hf_dir, local_files_only=True, trust_remote_code=False) + token_ids = tokenizer(OFT_ACTION_TOKEN, add_special_tokens=False)["input_ids"] + except Exception as exc: + raise StarVLAError(f"failed to load the pinned Qwen tokenizer from {hf_dir}: {exc}") from exc + if token_ids != [OFT_ACTION_TOKEN_ID]: + raise StarVLAError( + f"unexpected OFT action token mapping for {OFT_ACTION_TOKEN!r}: " + f"expected [{OFT_ACTION_TOKEN_ID}], got {token_ids}" + ) + return token_ids[0] + + +def normalization_metadata( + stats: dict[str, Any], action_dim: int, default_profile: str +) -> dict[str, Any]: + if default_profile not in stats: + raise StarVLAError( + f"default normalization profile {default_profile!r} is not present" + ) + profile_keys = [default_profile, *sorted(set(stats) - {default_profile})] + metadata: dict[str, Any] = { + "starvla.normalization.profile_count": len(stats), + "starvla.normalization.profile_keys": profile_keys, + "starvla.normalization.clip_actions": False, + "starvla.normalization.binary_threshold": 0.5, + "starvla.normalization.binary_comparison": "gt", + } + for index, key in enumerate(profile_keys): + profile = stats[key] + action = profile.get("action") + if not isinstance(action, dict): + raise StarVLAError(f"normalization profile {key!r} has no action object") + for field in ("q01", "q99", "mask"): + values = action.get(field) + if not isinstance(values, list) or len(values) != action_dim: + raise StarVLAError( + f"normalization profile {key!r} action.{field} must have {action_dim} values" + ) + metadata[f"starvla.normalization.profile.{index}.action_{field}"] = values + q01 = action["q01"] + q99 = action["q99"] + mask = action["mask"] + expected_mask = [True] * (action_dim - 1) + [False] + if any(type(value) is not bool for value in mask) or mask != expected_mask: + raise StarVLAError( + f"normalization profile {key!r} action.mask must be {expected_mask}, got {mask}" + ) + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*q01, *q99] + ): + raise StarVLAError(f"normalization profile {key!r} action quantiles must be finite numbers") + if any(q99[index] < q01[index] for index in range(action_dim - 1)): + raise StarVLAError(f"normalization profile {key!r} has q99 below q01") + metadata[f"starvla.normalization.profile.{index}.key"] = key + + state = profile.get("state") + if state is not None: + if not isinstance(state, dict): + raise StarVLAError(f"normalization profile {key!r} state must be an object") + state_q01 = state.get("q01") + state_q99 = state.get("q99") + if ( + not isinstance(state_q01, list) + or not isinstance(state_q99, list) + or not state_q01 + or len(state_q01) != len(state_q99) + ): + raise StarVLAError(f"normalization profile {key!r} has inconsistent state q01/q99") + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + for value in [*state_q01, *state_q99] + ): + raise StarVLAError(f"normalization profile {key!r} state quantiles must be finite numbers") + if any(upper < lower for lower, upper in zip(state_q01, state_q99)): + raise StarVLAError(f"normalization profile {key!r} state has q99 below q01") + metadata[f"starvla.normalization.profile.{index}.state_dimension"] = len(state_q01) + metadata[f"starvla.normalization.profile.{index}.state_q01"] = state_q01 + metadata[f"starvla.normalization.profile.{index}.state_q99"] = state_q99 + return metadata + + +def build_oft_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + action_token_id: int, + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + backbone = str(variant["backbone"]) + config = load_oft_config(policy_dir, surgery_manifest, backbone) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + datasets = config.get("datasets", {}) + vla_config = datasets.get("vla_data", {}) + action_horizon = int(action_config.get("action_horizon", int(action_config.get("future_action_window_size", 15)) + 1)) + if action_horizon != 16: + raise StarVLAError(f"unsupported OFT action horizon: {action_horizon}") + if dimensions["action_dim"] != 7: + raise StarVLAError(f"unsupported OFT action dimension: {dimensions['action_dim']}") + expected_dimensions = { + "qwen3_vl": (2560, 5120), + "qwen2_5_vl": (2048, 4096), + }.get(backbone) + if expected_dimensions is None: + raise StarVLAError(f"unsupported OFT Qwen backbone: {backbone!r}") + if ( + dimensions["input_dim"], + dimensions["hidden_dim"], + ) != expected_dimensions: + raise StarVLAError(f"unsupported OFT MLP dimensions: {dimensions}") + + action_tokens = OFT_ACTION_TOKEN * action_horizon + action_suffix = f" Please predict the next {action_horizon} robot actions: {action_tokens}." + image_size = vla_config.get("image_size", [224, 224]) + image_names = vla_config.get("obs", ["image_0"]) + if image_size != [224, 224] or not isinstance(image_names, list): + raise StarVLAError(f"unsupported OFT image config: image_size={image_size}, obs={image_names}") + + qwen = ( + _validate_pinned_qwen3vl_contract(hf_dir) + if backbone == "qwen3_vl" + else _validate_pinned_qwen25vl_contract(hf_dir) + ) + if qwen.get("hidden_size", dimensions["input_dim"]) != dimensions["input_dim"]: + raise StarVLAError( + "OFT policy input dimension does not match the staged Qwen backbone" + ) + + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": ( + "StarVLA Qwen3-VL OFT policy" + if backbone == "qwen3_vl" + else "StarVLA Qwen2.5-VL OFT policy" + ), + "starvla.schema_version": 1, + "starvla.framework": "oft", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": backbone, + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["input_dim"], + "starvla.qwen.input_embedding_size": ( + dimensions["input_dim"] * 4 + if backbone == "qwen3_vl" + else dimensions["input_dim"] + ), + "starvla.qwen.vocab_size": qwen.get("vocab_size", 151936), + "starvla.prompt.action_token": OFT_ACTION_TOKEN, + "starvla.prompt.action_token_id": action_token_id, + "starvla.prompt.action_suffix": action_suffix, + "starvla.prompt.cot_template": str(vla_config.get("CoT_prompt", "")), + "starvla.prompt.cot_enabled": bool(vla_config.get("CoT_prompt", "")), + "starvla.prompt.state_bins": 256, + "starvla.prompt.state_bin_min": -1.0, + "starvla.prompt.state_bin_max": 1.0, + "starvla.prompt.state_clip": False, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": action_horizon, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.oft.hidden_size": dimensions["hidden_dim"], + "starvla.oft.block_count": 2, + "starvla.oft.layer_norm_epsilon": OFT_LAYER_NORM_EPS, + } + if backbone == "qwen3_vl": + metadata.update(_runtime_image_metadata( + build_qwen3vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="OFT", + ) + )) + else: + metadata.update(_runtime_image_metadata( + build_qwen25vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="OFT", + ) + )) + stats = _load_json(policy_dir / "dataset_statistics.json") + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] + ) + ) + return metadata + + +def _validate_pinned_qwen3vl_contract(hf_dir: Path) -> dict[str, Any]: + qwen_config = _load_json(hf_dir / "config.json") + text_config = qwen_config.get("text_config", {}) + vision_config = qwen_config.get("vision_config", {}) + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + actual = { + "architecture": qwen_config.get("architectures"), + "vocab_size": text_config.get("vocab_size"), + "hidden_size": text_config.get("hidden_size"), + "layer_count": text_config.get("num_hidden_layers"), + "head_count": text_config.get("num_attention_heads"), + "head_count_kv": text_config.get("num_key_value_heads"), + "head_dim": text_config.get("head_dim"), + "vision_hidden_size": vision_config.get("hidden_size"), + "vision_layer_count": vision_config.get("depth"), + "vision_head_count": vision_config.get("num_heads"), + "vision_patch_size": vision_config.get("patch_size"), + "vision_temporal_patch_size": vision_config.get("temporal_patch_size"), + "vision_merge_size": vision_config.get("spatial_merge_size"), + "vision_deepstack": vision_config.get("deepstack_visual_indexes"), + "processor_size": preprocessor.get("size"), + "processor_patch_size": preprocessor.get("patch_size"), + "processor_temporal_patch_size": preprocessor.get("temporal_patch_size"), + "processor_merge_size": preprocessor.get("merge_size"), + "processor_class": preprocessor.get("processor_class"), + "image_processor_type": preprocessor.get("image_processor_type"), + "image_mean": preprocessor.get("image_mean"), + "image_std": preprocessor.get("image_std"), + } + expected = { + "architecture": ["Qwen3VLForConditionalGeneration"], + "vocab_size": 151936, + "hidden_size": 2560, + "layer_count": 36, + "head_count": 32, + "head_count_kv": 8, + "head_dim": 128, + "vision_hidden_size": 1024, + "vision_layer_count": 24, + "vision_head_count": 16, + "vision_patch_size": 16, + "vision_temporal_patch_size": 2, + "vision_merge_size": 2, + "vision_deepstack": [5, 11, 17], + "processor_size": { + "shortest_edge": QWEN3VL_PROCESSOR_MIN_PIXELS, + "longest_edge": QWEN3VL_PROCESSOR_MAX_PIXELS, + }, + "processor_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen3VLProcessor", + "image_processor_type": "Qwen2VLImageProcessorFast", + "image_mean": QWEN3VL_IMAGE_MEAN, + "image_std": QWEN3VL_IMAGE_STD, + } + if actual != expected: + raise StarVLAError(f"unexpected pinned Qwen config/processor contract: {actual}") + chat_template_path = hf_dir / "chat_template.json" + if not chat_template_path.is_file(): + raise StarVLAError(f"missing pinned Qwen chat template: {chat_template_path}") + return { + **actual, + "chat_template_sha256": sha256_file(chat_template_path), + } + + +def _validate_pinned_qwen25vl_contract(hf_dir: Path) -> dict[str, Any]: + qwen_config = _load_json(hf_dir / "config.json") + text_config = qwen_config.get("text_config") + if not isinstance(text_config, dict): + text_config = qwen_config + vision_config = qwen_config.get("vision_config", {}) + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + hidden_size = text_config.get("hidden_size") + head_count = text_config.get("num_attention_heads") + head_dim = ( + hidden_size // head_count + if isinstance(hidden_size, int) + and isinstance(head_count, int) + and head_count > 0 + and hidden_size % head_count == 0 + else None + ) + actual = { + "architecture": qwen_config.get("architectures"), + "model_type": qwen_config.get("model_type"), + "tie_word_embeddings": text_config.get("tie_word_embeddings"), + "vocab_size": text_config.get("vocab_size"), + "hidden_size": hidden_size, + "layer_count": text_config.get("num_hidden_layers"), + "head_count": head_count, + "head_count_kv": text_config.get("num_key_value_heads"), + "head_dim": head_dim, + "vision_hidden_size": vision_config.get("hidden_size"), + "vision_layer_count": vision_config.get("depth"), + "vision_head_count": vision_config.get("num_heads"), + "vision_patch_size": vision_config.get("patch_size"), + "vision_temporal_patch_size": vision_config.get("temporal_patch_size"), + "vision_merge_size": vision_config.get("spatial_merge_size"), + "vision_window_size": vision_config.get("window_size"), + "vision_full_attention_blocks": vision_config.get("fullatt_block_indexes"), + "vision_deepstack": [], + "processor_min_pixels": preprocessor.get("min_pixels"), + "processor_max_pixels": preprocessor.get("max_pixels"), + "processor_patch_size": preprocessor.get("patch_size"), + "processor_temporal_patch_size": preprocessor.get("temporal_patch_size"), + "processor_merge_size": preprocessor.get("merge_size"), + "processor_class": preprocessor.get("processor_class"), + "image_processor_type": preprocessor.get("image_processor_type"), + "image_mean": preprocessor.get("image_mean"), + "image_std": preprocessor.get("image_std"), + } + expected_image_processor_type = { + 151_936: "Qwen2VLImageProcessor", + 153_713: "Qwen2VLImageProcessorFast", + }.get(actual["vocab_size"]) + expected = { + "architecture": ["Qwen2_5_VLForConditionalGeneration"], + "model_type": "qwen2_5_vl", + "tie_word_embeddings": False, + "vocab_size": actual["vocab_size"], + "hidden_size": 2048, + "layer_count": 36, + "head_count": 16, + "head_count_kv": 2, + "head_dim": 128, + "vision_hidden_size": 1280, + "vision_layer_count": 32, + "vision_head_count": 16, + "vision_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "vision_window_size": 112, + "vision_full_attention_blocks": [7, 15, 23, 31], + "vision_deepstack": [], + "processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen2_5_VLProcessor", + "image_processor_type": expected_image_processor_type, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + } + if expected_image_processor_type is None: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL vocabulary size: {actual['vocab_size']!r}" + ) + if actual != expected: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL config/processor contract: {actual}" + ) + chat_template_path = hf_dir / "chat_template.json" + if not chat_template_path.is_file(): + # The action-expanded checkpoint publishes the same template as Jinja. + chat_template_path = hf_dir / "chat_template.jinja" + if not chat_template_path.is_file(): + raise StarVLAError(f"missing pinned Qwen2.5-VL chat template in {hf_dir}") + return { + **actual, + "chat_template_sha256": sha256_file(chat_template_path), + } + + +def _validate_pinned_qwenvl_contract( + hf_dir: Path, backbone: str +) -> dict[str, Any]: + if backbone == "qwen3_vl": + return _validate_pinned_qwen3vl_contract(hf_dir) + if backbone == "qwen2_5_vl": + return _validate_pinned_qwen25vl_contract(hf_dir) + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + +def _validate_image_config( + vla_config: dict[str, Any], + variant_label: str, + config_label: str, +) -> None: + if not isinstance(vla_config, dict): + raise StarVLAError(f"{variant_label} {config_label} vla_data must be an object") + if "obs_image_size" in vla_config: + raise StarVLAError( + f"{variant_label} does not support datasets.vla_data.obs_image_size" + ) + + +def build_qwen3vl_image_metadata( + vla_config: dict[str, Any], + qwen: dict[str, Any], + image_names: list[str], + *, + variant_label: str, + config_label: str = "effective config", +) -> dict[str, Any]: + """Build the Qwen3-VL image preprocessing contract.""" + _validate_image_config(vla_config, variant_label, config_label) + if not image_names or any(not isinstance(name, str) or not name for name in image_names): + raise StarVLAError(f"{variant_label} image names must be non-empty strings") + + processor_size = qwen.get("processor_size") + actual = { + "min_pixels": processor_size.get("shortest_edge") if isinstance(processor_size, dict) else None, + "max_pixels": processor_size.get("longest_edge") if isinstance(processor_size, dict) else None, + "processor_patch_size": qwen.get("processor_patch_size"), + "processor_temporal_patch_size": qwen.get("processor_temporal_patch_size"), + "processor_merge_size": qwen.get("processor_merge_size"), + "processor_class": qwen.get("image_processor_type"), + "image_mean": qwen.get("image_mean"), + "image_std": qwen.get("image_std"), + "vision_patch_size": qwen.get("vision_patch_size"), + "vision_temporal_patch_size": qwen.get("vision_temporal_patch_size"), + "vision_merge_size": qwen.get("vision_merge_size"), + } + expected = { + "min_pixels": QWEN3VL_PROCESSOR_MIN_PIXELS, + "max_pixels": QWEN3VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + "processor_class": "Qwen2VLImageProcessorFast", + "image_mean": QWEN3VL_IMAGE_MEAN, + "image_std": QWEN3VL_IMAGE_STD, + "vision_patch_size": QWEN3VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN3VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN3VL_SPATIAL_MERGE_SIZE, + } + if actual != expected: + raise StarVLAError(f"unexpected pinned Qwen dynamic image contract: {actual}") + + token_area = QWEN3VL_IMAGE_PATCH_SIZE**2 * QWEN3VL_SPATIAL_MERGE_SIZE**2 + if ( + QWEN3VL_PROCESSOR_MIN_PIXELS // token_area != QWEN3VL_MIN_IMAGE_TOKENS + or QWEN3VL_PROCESSOR_MAX_PIXELS // token_area != QWEN3VL_MAX_IMAGE_TOKENS + or QWEN3VL_PROCESSOR_MIN_PIXELS % token_area + or QWEN3VL_PROCESSOR_MAX_PIXELS % token_area + ): + raise StarVLAError("internal Qwen3-VL smart-resize image-token bounds drift") + metadata = { + key: list(value) if isinstance(value, list) else value + for key, value in QWEN3VL_DYNAMIC_IMAGE_METADATA.items() + } + metadata["starvla.image.count"] = len(image_names) + metadata["starvla.image.names"] = list(image_names) + return metadata + + +def build_qwen25vl_image_metadata( + vla_config: dict[str, Any], + qwen: dict[str, Any], + image_names: list[str], + *, + variant_label: str, + config_label: str = "effective config", +) -> dict[str, Any]: + """Build the Transformers 4.57 fast Qwen2.5-VL image contract.""" + _validate_image_config( + vla_config, variant_label, config_label + ) + if not image_names or any(not isinstance(name, str) or not name for name in image_names): + raise StarVLAError(f"{variant_label} image names must be non-empty strings") + expected = { + "processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "processor_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "processor_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + "vision_patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "vision_temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "vision_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + } + actual = {key: qwen.get(key) for key in expected} + if actual != expected: + raise StarVLAError( + f"unexpected pinned Qwen2.5-VL dynamic image contract: {actual}" + ) + + token_area = ( + QWEN25VL_IMAGE_PATCH_SIZE**2 * QWEN25VL_SPATIAL_MERGE_SIZE**2 + ) + if ( + QWEN25VL_PROCESSOR_MIN_PIXELS // token_area + != QWEN25VL_MIN_IMAGE_TOKENS + or QWEN25VL_PROCESSOR_MAX_PIXELS // token_area + != QWEN25VL_MAX_IMAGE_TOKENS + or QWEN25VL_PROCESSOR_MIN_PIXELS % token_area + or QWEN25VL_PROCESSOR_MAX_PIXELS % token_area + ): + raise StarVLAError( + "internal Qwen2.5-VL smart-resize image-token bounds drift" + ) + return { + "starvla.image.count": len(image_names), + "starvla.image.names": list(image_names), + "starvla.image.preprocessing_mode": "qwen2_5vl_smart_resize", + "starvla.image.framework_inference_pre_resize": False, + "starvla.image.framework_inference_pre_resize_config_key": + "datasets.vla_data.obs_image_size", + "starvla.image.processor_min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "starvla.image.processor_max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "starvla.image.processor_class": "Qwen2VLImageProcessorFast", + "starvla.image.processor_reference_transformers_version": "4.57.0", + "starvla.image.processor_do_convert_rgb": True, + "starvla.image.processor_do_resize": True, + "starvla.image.processor_resize_resample": "bicubic", + "starvla.image.processor_resize_antialias": True, + "starvla.image.processor_do_rescale": True, + "starvla.image.processor_rescale_factor": 1.0 / 255.0, + "starvla.image.processor_do_normalize": True, + "starvla.image.processor_image_mean": list(QWEN25VL_IMAGE_MEAN), + "starvla.image.processor_image_std": list(QWEN25VL_IMAGE_STD), + "starvla.image.patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "starvla.image.temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "starvla.image.spatial_merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "starvla.image.token_count_mode": + "dynamic_grid_thw_after_spatial_merge", + "starvla.image.min_token_count": QWEN25VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN25VL_MAX_IMAGE_TOKENS, + } + + +def _runtime_image_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + keys = ( + "starvla.image.count", + "starvla.image.names", + "starvla.image.processor_min_pixels", + "starvla.image.processor_max_pixels", + "starvla.image.patch_size", + "starvla.image.spatial_merge_size", + "starvla.image.min_token_count", + "starvla.image.max_token_count", + ) + return {key: metadata[key] for key in keys} + + +def build_groot_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + """Build the executable contract for a Qwen-VL GR00T head.""" + backbone = str(variant["backbone"]) + config = load_groot_config(policy_dir, surgery_manifest, backbone) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + diffusion_config = action_config.get("diffusion_model_cfg", {}) + vla_config = config.get("datasets", {}).get("vla_data", {}) + qwen = _validate_pinned_qwenvl_contract(hf_dir, backbone) + + expected_dimensions = GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE.get(backbone) + if expected_dimensions is None: + raise StarVLAError(f"unsupported GR00T Qwen backbone: {backbone!r}") + if dimensions != expected_dimensions: + raise StarVLAError(f"unsupported GR00T tensor dimensions: {dimensions}") + if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: + raise StarVLAError( + "GR00T cross-attention dimension does not match the staged Qwen backbone" + ) + + action_horizon = int( + action_config.get( + "action_horizon", + int(action_config.get("future_action_window_size", 15)) + 1, + ) + ) + if action_horizon != 16: + raise StarVLAError(f"unsupported GR00T action horizon: {action_horizon}") + head_count = int(diffusion_config.get("num_attention_heads", 0)) + head_dim = int(diffusion_config.get("attention_head_dim", 0)) + if head_count * head_dim != dimensions["dit_width"]: + raise StarVLAError("GR00T attention config does not match checkpoint shapes") + image_names = vla_config.get("obs", ["image_0"]) + if vla_config.get("image_size", [224, 224]) != [224, 224] or not isinstance(image_names, list): + raise StarVLAError("GR00T requires a 224x224 image configuration") + if vla_config.get("include_state", False) not in (False, "False"): + raise StarVLAError("GR00T state input is not supported") + + cot_template = str(vla_config.get("CoT_prompt", "")) + if not cot_template: + raise StarVLAError("GR00T CoT_prompt must not be empty") + + num_steps = int(action_config["num_inference_timesteps"]) + timestep_buckets = int(action_config["num_timestep_buckets"]) + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("GR00T timestep counts must be positive") + timestep_ids = [step * timestep_buckets // num_steps for step in range(num_steps)] + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": ( + "StarVLA Qwen3-VL GR00T policy" + if backbone == "qwen3_vl" + else "StarVLA Qwen2.5-VL GR00T policy" + ), + "starvla.schema_version": 1, + "starvla.framework": "groot", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": backbone, + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": ( + dimensions["qwen_hidden_dim"] * 4 + if backbone == "qwen3_vl" + else dimensions["qwen_hidden_dim"] + ), + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": cot_template, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": action_horizon, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.groot.dit_width": dimensions["dit_width"], + "starvla.groot.block_count": dimensions["block_count"], + "starvla.groot.attention_head_count": head_count, + "starvla.groot.attention_head_dim": head_dim, + "starvla.groot.cross_attention_dim": dimensions["qwen_hidden_dim"], + "starvla.groot.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.groot.ada_norm_epsilon": GROOT_DIT_NORM_EPS, + "starvla.groot.output_norm_epsilon": GROOT_OUTPUT_NORM_EPS, + "starvla.groot.output_dimension": dimensions["output_dim"], + "starvla.groot.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.groot.future_token_count": dimensions["future_token_count"], + "starvla.groot.action_position_count": dimensions["max_sequence_length"], + "starvla.groot.no_state_sequence_length": dimensions["future_token_count"] + action_horizon, + "starvla.groot.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.groot.timestep_ids": timestep_ids, + "starvla.groot.euler_dt": 1.0 / num_steps, + } + if backbone == "qwen3_vl": + metadata.update(_runtime_image_metadata( + build_qwen3vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="GR00T", + ) + )) + else: + metadata.update(_runtime_image_metadata( + build_qwen25vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="GR00T", + ) + )) + stats = _load_json(policy_dir / "dataset_statistics.json") + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] + ) + ) + return metadata + + +def build_pi_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + """Build the Qwen2.5-VL legacy PI executable contract.""" + if variant.get("framework") != "pi" or variant.get("backbone") != "qwen2_5_vl": + raise StarVLAError("legacy PI metadata requires the qwen25_pi catalog variant") + config = load_pi_config(policy_dir, surgery_manifest, str(variant["backbone"])) + framework = config.get("framework", {}) + action_config = framework.get("action_model", {}) + diffusion_config = action_config.get("diffusion_model_cfg", {}) + vla_config = config.get("datasets", {}).get("vla_data", {}) + qwen = _validate_pinned_qwen25vl_contract(hf_dir) + + if dimensions != PI_SUPPORTED_DIMENSIONS: + raise StarVLAError(f"unsupported legacy PI tensor dimensions: {dimensions}") + if qwen["hidden_size"] != dimensions["qwen_hidden_dim"]: + raise StarVLAError( + "legacy PI cross-attention dimension does not match the staged Qwen backbone" + ) + + action_horizon = int( + action_config.get( + "action_horizon", + int(action_config.get("future_action_window_size", 15)) + 1, + ) + ) + if action_horizon != 16: + raise StarVLAError(f"unsupported PI action horizon: {action_horizon}") + head_count = int(diffusion_config.get("num_attention_heads", 0)) + head_dim = int(diffusion_config.get("attention_head_dim", 0)) + if head_count * head_dim != dimensions["dit_width"]: + raise StarVLAError("PI attention config does not match checkpoint shapes") + cot_template = str(vla_config.get("CoT_prompt", "")) + image_names = vla_config.get("obs", ["image_0"]) + image_size = vla_config.get("image_size", [224, 224]) + if not cot_template: + raise StarVLAError("PI CoT_prompt must not be empty") + if ( + not isinstance(image_names, list) + or not isinstance(image_size, list) + or len(image_size) != 2 + or any(type(value) is not int or value <= 0 for value in image_size) + ): + raise StarVLAError("PI image configuration is invalid") + + num_steps = int(action_config["num_inference_timesteps"]) + timestep_buckets = int(action_config["num_timestep_buckets"]) + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("PI timestep counts must be positive") + continuous_times = [step / float(num_steps) for step in range(num_steps)] + timestep_ids = [int(value * timestep_buckets) for value in continuous_times] + hidden_tuple_indices = list( + range(qwen["layer_count"] - PI_BLOCK_COUNT + 1, qwen["layer_count"] + 1) + ) + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen2.5-VL legacy PI policy", + "general.source.uuid": surgery_manifest["bundle_uuid"], + "starvla.schema_version": 1, + "starvla.framework": "pi", + "starvla.model_type": variant["model_type"], + "starvla.backbone.arch": "qwen2_5_vl", + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.layer_count": qwen["layer_count"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": cot_template, + "starvla.conditioning.hidden_tuple_indices": hidden_tuple_indices, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": action_horizon, + "starvla.action.continuous_dimensions": [0, 1, 2, 3, 4, 5], + "starvla.action.binary_dimensions": [6], + "starvla.state.dimension": dimensions["state_dim"], + "starvla.pi.dit_width": dimensions["dit_width"], + "starvla.pi.block_count": dimensions["block_count"], + "starvla.pi.attention_head_count": head_count, + "starvla.pi.attention_head_dim": head_dim, + "starvla.pi.cross_attention_dim": dimensions["qwen_hidden_dim"], + "starvla.pi.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.pi.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.pi.state_token_count": 1, + "starvla.pi.future_token_count": dimensions["future_token_count"], + "starvla.pi.action_position_count": dimensions["max_sequence_length"], + "starvla.pi.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.pi.num_inference_timesteps": num_steps, + "starvla.pi.timestep_ids": timestep_ids, + "starvla.pi.euler_dt": 1.0 / num_steps, + "starvla.pi.ada_norm_epsilon": PI_DIT_NORM_EPS, + } + image_metadata = build_qwen25vl_image_metadata( + vla_config, + qwen, + image_names, + variant_label="legacy PI", + ) + image_metadata.update( + { + "starvla.image.framework_inference_pre_resize": True, + "starvla.image.framework_inference_pre_resize_config_key": + "datasets.vla_data.image_size", + "starvla.image.framework_inference_pre_resize_width": image_size[1], + "starvla.image.framework_inference_pre_resize_height": image_size[0], + } + ) + for key in ( + "starvla.image.count", + "starvla.image.names", + "starvla.image.processor_min_pixels", + "starvla.image.processor_max_pixels", + "starvla.image.patch_size", + "starvla.image.spatial_merge_size", + "starvla.image.min_token_count", + "starvla.image.max_token_count", + "starvla.image.framework_inference_pre_resize_width", + "starvla.image.framework_inference_pre_resize_height", + ): + metadata[key] = image_metadata[key] + + stats = _load_json(policy_dir / "dataset_statistics.json") + metadata.update( + normalization_metadata( + stats, dimensions["action_dim"], variant["default_unnorm_key"] + ) + ) + metadata["starvla.normalization.clip_actions"] = True + metadata["starvla.normalization.binary_comparison"] = "ge" + return metadata + + +def build_pi_v3_metadata( + policy_dir: Path, + hf_dir: Path, + variant: dict[str, Any], + surgery_manifest: dict[str, Any], + dimensions: dict[str, int], + text_filename: str, + mmproj_filename: str, +) -> dict[str, Any]: + config = load_pi_v3_config( + policy_dir, surgery_manifest, str(variant["backbone"]) + ) + framework = config.get("framework", {}) + action = framework.get("action_model", {}) + diffusion = action.get("diffusion_model_cfg", {}) + vla = config.get("datasets", {}).get("vla_data", {}) + image_names = vla.get("obs", ["image_0"]) + qwen = _validate_pinned_qwen3vl_contract(hf_dir) + + expected_dimensions = { + "qwen_hidden_dim": qwen.get("hidden_size"), + "dit_width": diffusion.get("action_dit_hidden_dim"), + "action_dim": action.get("action_dim"), + "block_count": diffusion.get("num_layers"), + } + if any(dimensions[key] != value for key, value in expected_dimensions.items()): + raise StarVLAError("PI-v3 config does not match the checkpoint tensor shapes") + if not isinstance(image_names, list) or not image_names or any( + not isinstance(name, str) or not name for name in image_names + ): + raise StarVLAError("PI-v3 config does not define observation image names") + + horizon = int(action["action_horizon"]) + num_steps = int(action["num_inference_timesteps"]) + timestep_buckets = int(action["num_timestep_buckets"]) + if horizon != 16: + raise StarVLAError(f"unsupported PI-v3 action horizon: {horizon}") + if num_steps <= 0 or timestep_buckets <= 0: + raise StarVLAError("PI-v3 timestep counts must be positive") + if not str(vla.get("CoT_prompt", "")): + raise StarVLAError("PI-v3 CoT_prompt must not be empty") + processor_size = qwen["processor_size"] + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen3-VL PI-v3 policy", + "general.source.uuid": surgery_manifest["bundle_uuid"], + "starvla.schema_version": 1, + "starvla.framework": "pi_v3", + "starvla.model_type": "starvla", + "starvla.backbone.arch": "qwen3_vl", + "starvla.bundle.uuid": surgery_manifest["bundle_uuid"], + "starvla.component.text.filename": text_filename, + "starvla.component.mmproj.filename": mmproj_filename, + "starvla.qwen.hidden_size": dimensions["qwen_hidden_dim"], + "starvla.qwen.input_embedding_size": 4 * dimensions["qwen_hidden_dim"], + "starvla.qwen.layer_count": qwen["layer_count"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": str(vla.get("CoT_prompt", "")), + "starvla.image.count": len(image_names), + "starvla.image.names": image_names, + "starvla.image.processor_min_pixels": processor_size["shortest_edge"], + "starvla.image.processor_max_pixels": processor_size["longest_edge"], + "starvla.image.patch_size": qwen["processor_patch_size"], + "starvla.image.spatial_merge_size": qwen["processor_merge_size"], + "starvla.image.min_token_count": QWEN3VL_MIN_IMAGE_TOKENS, + "starvla.image.max_token_count": QWEN3VL_MAX_IMAGE_TOKENS, + "starvla.action.dimension": dimensions["action_dim"], + "starvla.action.horizon": horizon, + "starvla.action.continuous_dimensions": list(range(dimensions["action_dim"] - 1)), + "starvla.action.binary_dimensions": [dimensions["action_dim"] - 1], + "starvla.pi_v3.dit_width": dimensions["dit_width"], + "starvla.pi_v3.block_count": dimensions["block_count"], + "starvla.pi_v3.projector_count": dimensions["projector_count"], + "starvla.pi_v3.attention_head_count": diffusion["num_attention_heads"], + "starvla.pi_v3.attention_head_dim": diffusion["attention_head_dim"], + "starvla.pi_v3.feed_forward_dim": dimensions["feed_forward_dim"], + "starvla.pi_v3.mlp_hidden_dimension": dimensions["mlp_hidden_dim"], + "starvla.pi_v3.future_token_count": dimensions["future_token_count"], + "starvla.pi_v3.action_position_count": dimensions["max_sequence_length"], + "starvla.pi_v3.no_state_sequence_length": dimensions["future_token_count"] + horizon, + "starvla.pi_v3.timestep_projection_dim": dimensions["timestep_dim"], + "starvla.pi_v3.num_timestep_buckets": timestep_buckets, + "starvla.pi_v3.num_inference_timesteps": num_steps, + "starvla.pi_v3.ada_norm_epsilon": PI_V3_DIT_NORM_EPS, + "starvla.pi_v3.projector_norm_epsilon": PI_V3_PROJECTOR_NORM_EPS, + "starvla.pi_v3.euler_dt": 1.0 / num_steps, + } + metadata.update( + normalization_metadata( + _load_json(policy_dir / "dataset_statistics.json"), + dimensions["action_dim"], + variant["default_unnorm_key"], + ) + ) + return metadata + + +def convert_policy( + policy_dir: Path, + hf_dir: Path, + surgery_manifest_path: Path, + output: Path, + catalog_path: Path, + dtype: str, + text_filename: str, + mmproj_filename: str, +) -> None: + catalog = load_catalog(catalog_path) + surgery_manifest = _load_json(surgery_manifest_path) + variant = get_variant(catalog, str(surgery_manifest.get("variant", ""))) + framework = str(variant["framework"]) + validators = { + "oft": validate_oft_tensors, + "groot": validate_groot_tensors, + "pi": validate_pi_tensors, + "pi_v3": validate_pi_v3_tensors, + } + tensor_maps = { + "oft": OFT_TENSOR_MAP, + "groot": GROOT_TENSOR_MAP, + "pi": PI_TENSOR_MAP, + "pi_v3": PI_V3_TENSOR_MAP, + } + if framework not in validators: + raise StarVLAError( + f"surgery variant {surgery_manifest.get('variant')!r} has no policy converter" + ) + + validate_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets( + hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen" + ) + verify_staged_assets( + policy_dir, surgery_manifest.get("policy_assets", {}), component="policy" + ) + tensors = load_policy_tensors(policy_dir) + dimensions = validators[framework](tensors) + metadata_args = ( + policy_dir, + hf_dir, + variant, + surgery_manifest, + dimensions, + ) + if framework == "oft": + metadata = build_oft_metadata( + *metadata_args, + resolve_action_token_id(hf_dir), + text_filename, + mmproj_filename, + ) + else: + metadata_builder = { + "groot": build_groot_metadata, + "pi": build_pi_metadata, + "pi_v3": build_pi_v3_metadata, + }[framework] + metadata = metadata_builder( + *metadata_args, + text_filename, + mmproj_filename, + ) + + writer_dir = Path(__file__).resolve().parents[1] / "pi0" + sys.path.insert(0, str(writer_dir)) + try: + from gguf_writer import write_gguf_arrays + except ImportError as exc: + raise StarVLAError( + f"failed to import repository GGUF writer adapter: {exc}" + ) from exc + + def arrays(): + for source_name, destination_name in tensor_maps[framework].items(): + tensor = tensors[source_name] + array = tensor.detach().float().cpu().numpy() + yield ( + destination_name, + [int(dimension) for dimension in tensor.shape], + np.asarray(array), + dtype, + ) + + _write_gguf_arrays_no_overwrite( + output, metadata, arrays(), write_gguf_arrays + ) + + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + default="oft", + choices=( + "oft", + "groot", + "pi_v3", + "qwen25_oft", + "qwen25_groot", + "qwen25_pi", + ), + ) + parser.add_argument("--policy-dir", type=Path, required=True) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--dtype", + choices=("fp32", "f16", "bf16"), + default=DEFAULT_POLICY_DTYPE, + ) + parser.add_argument("--text-filename") + parser.add_argument("--mmproj-filename") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if args.output.exists() or args.output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {args.output}") + text_filename = args.text_filename or default_text_filename( + args.variant, DEFAULT_TEXT_DTYPE + ) + mmproj_filename = args.mmproj_filename or default_mmproj_filename( + args.variant, DEFAULT_MMPROJ_DTYPE + ) + convert_policy( + policy_dir=args.policy_dir, + hf_dir=args.hf_dir, + surgery_manifest_path=args.surgery_manifest, + output=args.output, + catalog_path=args.catalog, + dtype=args.dtype, + text_filename=text_filename, + mmproj_filename=mmproj_filename, + ) + print(f"policy GGUF: {args.output}") + return 0 + except (StarVLAError, OSError, json.JSONDecodeError, ValueError, TypeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py new file mode 100644 index 0000000..f9e4ef2 --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_qwen25_fast.py @@ -0,0 +1,1454 @@ +#!/usr/bin/env python3 +"""Stage and convert a Qwen2.5-VL StarVLA FAST checkpoint.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any, Mapping, Sequence + +import numpy as np + +from convert_starvla_policy_to_gguf import normalization_metadata +from convert_starvla_qwen_to_gguf import build_commands, verify_llama_checkout +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + atomic_write_json, + build_inventory, + get_qwen_asset, + get_variant, + inventory_summary, + load_catalog, + load_checkpoint_state, + bundle_uuid, + portable_source_record, + sha256_file, + staged_qwen_asset_hashes, + validate_qwen_vlm_destination_names, + verify_catalog_files, + verify_checkpoint_file, + verify_staged_assets, + verify_staged_shards, +) +from starvla_surgery import ( + copy_policy_assets, + copy_qwen_assets, + parse_size, + write_safetensor_shards, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_GGUF_PY = REPOSITORY_ROOT / "third_party" / "llama.cpp" / "gguf-py" +if not LLAMA_GGUF_PY.is_dir(): + raise ImportError( + "third_party/llama.cpp/gguf-py is required; initialize the llama.cpp submodule" + ) +sys.path.insert(0, str(LLAMA_GGUF_PY)) + +import gguf # noqa: E402 + + +VARIANT_KEY = "qwen25_fast" +QWEN_ASSET_KEY = "qwen2_5_vl_3b_instruct_action" +FAST_CODEC_ASSET_KEY = "fast_codec" + +MODEL_TYPE = "starvla" +BACKBONE = "qwen2_5_vl" +FRAMEWORK = "fast" + +ACTION_TOKEN_COUNT = 2048 +ACTION_TOKEN_MIN = 151665 +ACTION_TOKEN_MAX = 153712 +ACTION_DIM = 7 +ACTION_HORIZON = 16 +MAX_LENGTH = 2048 + +BOS_TOKEN_ID = 151643 +EOS_TOKEN_IDS = [151645, 151643] +PAD_TOKEN_ID = 151643 +GENERATION_CONTRACT = { + "max_length": MAX_LENGTH, + "do_sample": True, + "temperature": 0.1, + "top_k": 1, + "top_p": 0.001, + "repetition_penalty": 1.05, + "bos_token_id": BOS_TOKEN_ID, + "eos_token_id": EOS_TOKEN_IDS, + "pad_token_id": PAD_TOKEN_ID, +} + +EXPECTED_INVENTORY = { + "total_tensors": 825, + "vlm_tensors": 825, + "policy_tensors": 0, + "visual_tensors": 390, + "text_tensors": 434, + "lm_head_tensors": 1, + "total_numel": 4_073_066_496, + "vlm_numel": 4_073_066_496, + "policy_numel": 0, + "total_nbytes": 8_146_132_992, + "vlm_nbytes": 8_146_132_992, + "policy_nbytes": 0, + "dtypes": {"bfloat16": 825}, + "storage_alias_groups": 0, +} + +TEXT_FILENAME = "qwen-qwen25-fast-bf16.gguf" +MMPROJ_FILENAME = "mmproj-qwen25-fast-bf16.gguf" +POLICY_FILENAME = "policy-qwen25-fast.gguf" +STAGING_MANIFEST_FILENAME = "qwen25-fast-staging-manifest.json" +BUNDLE_MANIFEST_FILENAME = "conversion_manifest.json" + +ACTION_NAMES = ["x", "y", "z", "roll", "pitch", "yaw", "gripper"] + +ACTION_TOKEN_MAP_TENSOR = "starvla.policy.fast.action_token_map" +CODEC_TOKEN_OFFSETS_TENSOR = "starvla.policy.fast.codec.token_offsets" +CODEC_TOKEN_BYTES_TENSOR = "starvla.policy.fast.codec.token_bytes" +FAST_RUNTIME_TENSOR_NAMES = { + ACTION_TOKEN_MAP_TENSOR, + CODEC_TOKEN_OFFSETS_TENSOR, + CODEC_TOKEN_BYTES_TENSOR, +} +QWEN25VL_PROCESSOR_MIN_PIXELS = 3_136 +QWEN25VL_PROCESSOR_MAX_PIXELS = 12_845_056 +QWEN25VL_IMAGE_PATCH_SIZE = 14 +QWEN25VL_TEMPORAL_PATCH_SIZE = 2 +QWEN25VL_SPATIAL_MERGE_SIZE = 2 +QWEN25VL_MIN_IMAGE_TOKENS = 4 +QWEN25VL_MAX_IMAGE_TOKENS = 16_384 +QWEN25VL_IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] +QWEN25VL_IMAGE_STD = [0.26862954, 0.26130258, 0.27577711] + +def load_json_object(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def validate_catalog_contract( + catalog: Mapping[str, Any], +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + entry = get_variant(catalog, VARIANT_KEY) + expected = { + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "qwen_asset": QWEN_ASSET_KEY, + "policy_prefixes": [], + } + mismatches = [ + f"{key}: expected {value!r}, got {entry.get(key)!r}" + for key, value in expected.items() + if entry.get(key) != value + ] + if entry.get("policy_tensors") not in (None, []): + mismatches.append("policy_tensors: FAST must not split a separate policy head") + bundle_uuid(entry, catalog) + + qwen_name, qwen_entry = get_qwen_asset(catalog, entry) + if qwen_name != QWEN_ASSET_KEY: + mismatches.append(f"Qwen asset: expected {QWEN_ASSET_KEY!r}, got {qwen_name!r}") + codec_entry = catalog.get("shared_assets", {}).get(FAST_CODEC_ASSET_KEY) + if not isinstance(codec_entry, dict): + mismatches.append("FAST codec: missing pinned shared asset") + codec_entry = {} + if mismatches: + raise StarVLAError("Qwen2.5 FAST catalog contract mismatch: " + "; ".join(mismatches)) + return entry, qwen_entry, codec_entry + + +def validate_qwen_config(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "config.json") + text = config.get("text_config") + vision = config.get("vision_config") + if not isinstance(text, dict) or not isinstance(vision, dict): + raise StarVLAError("Qwen2.5 FAST config has no text_config/vision_config object") + expected_top = { + "architectures": ["Qwen2_5_VLForConditionalGeneration"], + "model_type": "qwen2_5_vl", + "dtype": "bfloat16", + "vocab_size": 151936, + "image_token_id": 151655, + "video_token_id": 151656, + "vision_token_id": 151654, + "vision_start_token_id": 151652, + "vision_end_token_id": 151653, + } + expected_text = { + "model_type": "qwen2_5_vl_text", + "dtype": "bfloat16", + "hidden_size": 2048, + "intermediate_size": 11008, + "num_hidden_layers": 36, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "vocab_size": 153713, + "tie_word_embeddings": True, + } + expected_vision = { + "depth": 32, + "hidden_size": 1280, + "intermediate_size": 3420, + "num_heads": 16, + "out_hidden_size": 2048, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "window_size": 112, + "fullatt_block_indexes": [7, 15, 23, 31], + } + mismatches = [] + for owner_name, owner, contract in ( + ("config", config, expected_top), + ("text_config", text, expected_text), + ("vision_config", vision, expected_vision), + ): + for key, value in contract.items(): + if owner.get(key) != value: + mismatches.append( + f"{owner_name}.{key}: expected {value!r}, got {owner.get(key)!r}" + ) + if mismatches: + raise StarVLAError("Qwen2.5 FAST config mismatch: " + "; ".join(mismatches)) + return { + "text_hidden_size": 2048, + "text_layers": 36, + "text_attention_heads": 16, + "text_key_value_heads": 2, + "text_attention_head_dim": 128, + "vision_hidden_size": 1280, + "vision_layers": 32, + "vision_attention_heads": 16, + "vision_full_attention_blocks": [7, 15, 23, 31], + "vision_deepstack": False, + "vocab_size": 153713, + } + + +def validate_qwen_processor(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "preprocessor_config.json") + expected = { + "do_convert_rgb": True, + "do_normalize": True, + "do_rescale": True, + "do_resize": True, + "image_mean": QWEN25VL_IMAGE_MEAN, + "image_std": QWEN25VL_IMAGE_STD, + "image_processor_type": "Qwen2VLImageProcessorFast", + "max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "processor_class": "Qwen2_5_VLProcessor", + "resample": 3, + "rescale_factor": 1.0 / 255.0, + "temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + } + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in expected.items() + if config.get(key) != value + ] + expected_size = { + "longest_edge": QWEN25VL_PROCESSOR_MAX_PIXELS, + "shortest_edge": QWEN25VL_PROCESSOR_MIN_PIXELS, + } + if config.get("size") != expected_size: + mismatches.append( + f"size: expected {expected_size!r}, got {config.get('size')!r}" + ) + chat_template = qwen_dir / "chat_template.jinja" + if not chat_template.is_file() or chat_template.stat().st_size == 0: + mismatches.append("chat_template.jinja: missing or empty") + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST processor contract mismatch: " + "; ".join(mismatches) + ) + return { + "min_pixels": QWEN25VL_PROCESSOR_MIN_PIXELS, + "max_pixels": QWEN25VL_PROCESSOR_MAX_PIXELS, + "patch_size": QWEN25VL_IMAGE_PATCH_SIZE, + "temporal_patch_size": QWEN25VL_TEMPORAL_PATCH_SIZE, + "merge_size": QWEN25VL_SPATIAL_MERGE_SIZE, + "min_image_tokens": QWEN25VL_MIN_IMAGE_TOKENS, + "max_image_tokens": QWEN25VL_MAX_IMAGE_TOKENS, + "image_mean": list(QWEN25VL_IMAGE_MEAN), + "image_std": list(QWEN25VL_IMAGE_STD), + "chat_template_sha256": sha256_file(chat_template), + } + + +def validate_generation_config(qwen_dir: Path) -> dict[str, Any]: + config = load_json_object(qwen_dir / "generation_config.json") + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in GENERATION_CONTRACT.items() + if key != "max_length" and config.get(key) != value + ] + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST generation_config mismatch: " + "; ".join(mismatches) + ) + return dict(GENERATION_CONTRACT) + + +def _action_mapping(path: Path) -> dict[str, int]: + raw = load_json_object(path) + mapping: dict[str, int] = {} + for key, value in raw.items(): + if not isinstance(value, int) or isinstance(value, bool): + raise StarVLAError(f"FAST action mapping has a non-integer ID for {key!r}") + mapping[key] = value + return mapping + + +def validate_action_token_mapping(qwen_dir: Path) -> dict[str, Any]: + expected = { + f"": ACTION_TOKEN_MIN + index + for index in range(ACTION_TOKEN_COUNT) + } + primary = _action_mapping(qwen_dir / "added_token_id_map.json") + added = { + key: value + for key, value in _action_mapping(qwen_dir / "added_tokens.json").items() + if key.startswith(" Qwen 151665..153712 mapping" + ) + if added != expected: + raise StarVLAError("added_tokens.json disagrees with the pinned FAST action mapping") + + tokenizer = load_json_object(qwen_dir / "tokenizer.json") + tokenizer_action = { + str(record.get("content")): record + for record in tokenizer.get("added_tokens", []) + if isinstance(record, dict) + and str(record.get("content", "")).startswith(" dict[str, Any]: + hashes = verify_catalog_files(codec_dir, codec_entry) + config = load_json_object(codec_dir / "processor_config.json") + expected = { + "processor_class": "UniversalActionProcessor", + "scale": 10, + "vocab_size": ACTION_TOKEN_COUNT, + "min_token": -354, + "action_dim": None, + "time_horizon": None, + } + mismatches = [ + f"{key}: expected {value!r}, got {config.get(key)!r}" + for key, value in expected.items() + if config.get(key) != value + ] + if mismatches: + raise StarVLAError("FAST codec config mismatch: " + "; ".join(mismatches)) + return { + "repo_id": codec_entry["repo_id"], + "revision": codec_entry["revision"], + "scale": 10, + "min_token": -354, + "vocab_size": ACTION_TOKEN_COUNT, + "action_dim": ACTION_DIM, + "time_horizon": ACTION_HORIZON, + "files": hashes, + } + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _byte_level_inverse_alphabet() -> dict[int, int]: + direct = { + *range(0x21, 0x7F), + *range(0xA1, 0xAD), + *range(0xAE, 0x100), + } + inverse = {value: value for value in direct} + extra = 0 + for value in range(0x100): + if value not in direct: + inverse[0x100 + extra] = value + extra += 1 + if extra != 68 or len(inverse) != 256: + raise AssertionError("internal GPT-2 ByteLevel alphabet construction drift") + return inverse + + +def compile_fast_runtime_tensors( + qwen_dir: Path, + codec_dir: Path, +) -> dict[str, np.ndarray]: + """Compile the pinned HF FAST decode assets into runtime-only integer tables.""" + processor = load_json_object(codec_dir / "processor_config.json") + expected_processor = { + "processor_class": "UniversalActionProcessor", + "scale": 10, + "vocab_size": ACTION_TOKEN_COUNT, + "min_token": -354, + "action_dim": None, + "time_horizon": None, + } + mismatches = [ + f"{key}: expected {value!r}, got {processor.get(key)!r}" + for key, value in expected_processor.items() + if processor.get(key) != value + ] + if mismatches: + raise StarVLAError( + "FAST runtime processor contract mismatch: " + "; ".join(mismatches) + ) + + tokenizer = load_json_object(codec_dir / "tokenizer.json") + decoder = tokenizer.get("decoder") + model = tokenizer.get("model") + if ( + tokenizer.get("version") != "1.0" + or tokenizer.get("added_tokens") != [] + or decoder + != { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": True, + "use_regex": True, + } + or not isinstance(model, dict) + or model.get("type") != "BPE" + or not isinstance(model.get("vocab"), dict) + ): + raise StarVLAError( + "FAST tokenizer is not the pinned ByteLevel BPE decode contract" + ) + vocab = model["vocab"] + if len(vocab) != ACTION_TOKEN_COUNT: + raise StarVLAError( + f"FAST tokenizer vocabulary must contain {ACTION_TOKEN_COUNT} entries" + ) + vocab_by_id: list[str | None] = [None] * ACTION_TOKEN_COUNT + for piece, token_id in vocab.items(): + if ( + not isinstance(piece, str) + or not piece + or isinstance(token_id, bool) + or not isinstance(token_id, int) + or token_id < 0 + or token_id >= ACTION_TOKEN_COUNT + or vocab_by_id[token_id] is not None + ): + raise StarVLAError("FAST tokenizer vocabulary IDs are not an exact bijection") + vocab_by_id[token_id] = piece + if any(piece is None for piece in vocab_by_id): + raise StarVLAError("FAST tokenizer vocabulary has missing IDs") + + inverse_alphabet = _byte_level_inverse_alphabet() + token_offsets = [0] + flattened = bytearray() + for token_id, optional_piece in enumerate(vocab_by_id): + if optional_piece is None: + raise AssertionError("FAST vocabulary completeness check failed") + for character in optional_piece: + byte_value = inverse_alphabet.get(ord(character)) + if byte_value is None: + raise StarVLAError( + "FAST tokenizer piece contains a code point outside the " + f"ByteLevel alphabet at token ID {token_id}" + ) + flattened.append(byte_value) + token_offsets.append(len(flattened)) + + validate_action_token_mapping(qwen_dir) + raw_mapping = _action_mapping(qwen_dir / "added_token_id_map.json") + action_token_map = np.asarray( + [ + raw_mapping[f""] + for token_id in range(ACTION_TOKEN_COUNT) + ], + dtype=np.int32, + ) + offsets = np.asarray(token_offsets, dtype=np.int32) + token_bytes = ( + np.frombuffer(bytes(flattened), dtype=np.uint8).view(np.int8).copy() + ) + + if ( + offsets.shape != (ACTION_TOKEN_COUNT + 1,) + or offsets[0] != 0 + or offsets[-1] != token_bytes.size + or np.any(np.diff(offsets) <= 0) + or action_token_map.tolist() + != list(range(ACTION_TOKEN_MIN, ACTION_TOKEN_MAX + 1)) + ): + raise StarVLAError("compiled FAST runtime tensor shape/content mismatch") + return { + ACTION_TOKEN_MAP_TENSOR: action_token_map, + CODEC_TOKEN_OFFSETS_TENSOR: offsets, + CODEC_TOKEN_BYTES_TENSOR: token_bytes, + } + + +def _normalize_gguf_metadata_value(value: Any) -> Any: + if isinstance(value, bool) or isinstance(value, str) or value is None: + return value + if isinstance(value, int): + if value < -(2**31) or value >= 2**31: + raise StarVLAError(f"GGUF int32 metadata value is out of range: {value}") + return value + if isinstance(value, float): + if not math.isfinite(value): + raise StarVLAError("GGUF metadata floats must be finite") + return float(np.float32(value)) + if isinstance(value, list): + if not value: + raise StarVLAError("runtime GGUF metadata arrays must not be empty") + return [_normalize_gguf_metadata_value(item) for item in value] + if isinstance(value, dict): + return { + str(key): _normalize_gguf_metadata_value(item) + for key, item in value.items() + } + raise StarVLAError(f"unsupported GGUF metadata value: {value!r}") + + +def build_fast_runtime_policy( + *, + manifest: Mapping[str, Any], + entry: Mapping[str, Any], + codec_entry: Mapping[str, Any], + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, +) -> tuple[dict[str, Any], dict[str, np.ndarray]]: + source = manifest.get("source") + bundle_uuid = manifest.get("bundle_uuid") + if not isinstance(source, Mapping) or not isinstance(bundle_uuid, str): + raise StarVLAError("FAST staging manifest lacks source/bundle provenance") + + qwen = validate_qwen_config(qwen_dir) + processor = validate_qwen_processor(qwen_dir) + generation = validate_generation_config(qwen_dir) + codec = validate_fast_codec(codec_dir, codec_entry) + effective = effective_fast_config(source_dir) + stats = load_json_object(source_dir / "dataset_statistics.json") + arrays = compile_fast_runtime_tensors(qwen_dir, codec_dir) + offsets = arrays[CODEC_TOKEN_OFFSETS_TENSOR] + token_bytes = arrays[CODEC_TOKEN_BYTES_TENSOR] + if ( + effective.get("image_count") != 1 + or effective.get("action_dim") != ACTION_DIM + or effective.get("action_horizon") != ACTION_HORIZON + ): + raise StarVLAError("effective FAST source contract is incompatible") + + metadata: dict[str, Any] = { + "general.architecture": "starvla-policy", + "general.name": "StarVLA Qwen2.5-VL FAST policy", + "general.source.uuid": bundle_uuid, + "starvla.schema_version": 1, + "starvla.framework": FRAMEWORK, + "starvla.model_type": MODEL_TYPE, + "starvla.backbone.arch": BACKBONE, + "starvla.bundle.uuid": bundle_uuid, + "starvla.component.text.filename": TEXT_FILENAME, + "starvla.component.mmproj.filename": MMPROJ_FILENAME, + "starvla.qwen.hidden_size": qwen["text_hidden_size"], + "starvla.qwen.input_embedding_size": qwen["text_hidden_size"], + "starvla.qwen.layer_count": qwen["text_layers"], + "starvla.qwen.vocab_size": qwen["vocab_size"], + "starvla.prompt.cot_template": effective["cot_prompt"], + "starvla.action.dimension": ACTION_DIM, + "starvla.action.horizon": ACTION_HORIZON, + "starvla.action.continuous_dimensions": list(range(ACTION_DIM - 1)), + "starvla.action.binary_dimensions": [ACTION_DIM - 1], + "starvla.image.count": effective["image_count"], + "starvla.image.names": effective["image_names"], + "starvla.image.processor_min_pixels": processor["min_pixels"], + "starvla.image.processor_max_pixels": processor["max_pixels"], + "starvla.image.patch_size": processor["patch_size"], + "starvla.image.spatial_merge_size": processor["merge_size"], + "starvla.image.min_token_count": processor["min_image_tokens"], + "starvla.image.max_token_count": processor["max_image_tokens"], + "starvla.fast.generation.max_length": generation["max_length"], + "starvla.fast.generation.eos_token_ids": generation["eos_token_id"], + "starvla.fast.generation.top_k": generation["top_k"], + "starvla.fast.generation.repetition_penalty": generation["repetition_penalty"], + "starvla.fast.action_token.count": ACTION_TOKEN_COUNT, + "starvla.fast.codec.scale": codec["scale"], + "starvla.fast.codec.min_token": codec["min_token"], + "starvla.fast.codec.vocab_size": ACTION_TOKEN_COUNT, + "starvla.fast.codec.time_horizon": ACTION_HORIZON, + "starvla.fast.codec.action_dimension": ACTION_DIM, + "starvla.fast.codec.token_offsets_count": int(offsets.size), + "starvla.fast.codec.token_bytes_count": int(token_bytes.size), + } + metadata.update( + normalization_metadata(stats, ACTION_DIM, str(entry["default_unnorm_key"])) + ) + return { + key: _normalize_gguf_metadata_value(value) + for key, value in metadata.items() + }, arrays + + +def _add_runtime_metadata(writer: Any, metadata: Mapping[str, Any]) -> None: + for key in sorted(metadata): + if key == "general.architecture": + continue + value = metadata[key] + if isinstance(value, str): + if not value: + raise StarVLAError(f"GGUF string metadata must be non-empty: {key}") + writer.add_string(key, value) + elif isinstance(value, bool): + writer.add_bool(key, value) + elif isinstance(value, int): + writer.add_int32(key, value) + elif isinstance(value, float): + writer.add_float32(key, value) + elif isinstance(value, list): + if not value: + raise StarVLAError(f"GGUF array metadata must be non-empty: {key}") + writer.add_array(key, value) + else: + raise StarVLAError(f"unsupported GGUF metadata value for {key}: {value!r}") + + +def write_fast_runtime_policy_gguf( + path: Path, + metadata: Mapping[str, Any], + arrays: Mapping[str, np.ndarray], +) -> None: + if path.exists(): + raise StarVLAError(f"refusing to overwrite runtime policy GGUF: {path}") + if set(arrays) != FAST_RUNTIME_TENSOR_NAMES: + raise StarVLAError("FAST runtime GGUF tensor set is incomplete") + path.parent.mkdir(parents=True, exist_ok=True) + writer = gguf.GGUFWriter( + path, + arch="starvla-policy", + use_temp_file=True, + ) + try: + _add_runtime_metadata(writer, metadata) + for name in ( + ACTION_TOKEN_MAP_TENSOR, + CODEC_TOKEN_OFFSETS_TENSOR, + CODEC_TOKEN_BYTES_TENSOR, + ): + array = np.ascontiguousarray(arrays[name]) + writer.add_tensor(name, array) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + except BaseException: + path.unlink(missing_ok=True) + raise + + +def _gguf_field(reader: Any, key: str) -> Any: + field = reader.get_field(key) + if field is None: + raise StarVLAError(f"FAST runtime GGUF is missing metadata: {key}") + return field.contents() + + +def _metadata_matches(actual: Any, expected: Any) -> bool: + if isinstance(expected, bool): + return isinstance(actual, bool) and actual is expected + if isinstance(expected, float): + return isinstance(actual, (int, float)) and math.isclose( + actual, + expected, + rel_tol=1e-7, + abs_tol=1e-7, + ) + if isinstance(expected, list): + return ( + isinstance(actual, list) + and len(actual) == len(expected) + and all( + _metadata_matches(actual_item, expected_item) + for actual_item, expected_item in zip(actual, expected) + ) + ) + if isinstance(expected, dict): + return ( + isinstance(actual, dict) + and set(actual) == set(expected) + and all( + _metadata_matches(actual[key], expected[key]) + for key in expected + ) + ) + return actual == expected + + +def validate_fast_runtime_policy_gguf( + path: Path, + *, + expected_metadata: Mapping[str, Any] | None = None, + expected_arrays: Mapping[str, np.ndarray] | None = None, +) -> dict[str, Any]: + if not path.is_file() or path.stat().st_size == 0: + raise StarVLAError(f"missing FAST policy GGUF: {path}") + try: + reader = gguf.GGUFReader(path) + except Exception as exc: + raise StarVLAError(f"failed to read FAST policy GGUF: {exc}") from exc + if _gguf_field(reader, "general.architecture") != "starvla-policy": + raise StarVLAError("FAST policy GGUF has the wrong architecture") + + vocab_size = int(_gguf_field(reader, "starvla.fast.codec.vocab_size")) + token_bytes_count = int( + _gguf_field(reader, "starvla.fast.codec.token_bytes_count") + ) + action_dim = int(_gguf_field(reader, "starvla.action.dimension")) + tensors = {tensor.name: tensor for tensor in reader.tensors} + expected_shapes = { + ACTION_TOKEN_MAP_TENSOR: ("I32", [vocab_size]), + CODEC_TOKEN_OFFSETS_TENSOR: ("I32", [vocab_size + 1]), + CODEC_TOKEN_BYTES_TENSOR: ("I8", [token_bytes_count]), + } + if set(tensors) != set(expected_shapes) or len(tensors) != len(reader.tensors): + raise StarVLAError(f"FAST policy tensor set mismatch: {sorted(tensors)}") + for name, (dtype, shape) in expected_shapes.items(): + tensor = tensors[name] + if tensor.tensor_type.name != dtype or list(map(int, tensor.shape)) != shape: + raise StarVLAError(f"FAST policy tensor shape/type mismatch: {name}") + + action_map = np.asarray(tensors[ACTION_TOKEN_MAP_TENSOR].data, dtype=np.int32).reshape(-1) + offsets = np.asarray(tensors[CODEC_TOKEN_OFFSETS_TENSOR].data, dtype=np.int32).reshape(-1) + token_bytes = np.asarray(tensors[CODEC_TOKEN_BYTES_TENSOR].data, dtype=np.int8).reshape(-1) + if ( + np.any(action_map < 0) + or np.unique(action_map).size != vocab_size + or offsets[0] != 0 + or offsets[-1] != token_bytes.size + or np.any(np.diff(offsets) <= 0) + ): + raise StarVLAError("FAST policy codec tensors are invalid") + + profile_count = int(_gguf_field(reader, "starvla.normalization.profile_count")) + profile_keys = _gguf_field(reader, "starvla.normalization.profile_keys") + if profile_count <= 0 or not isinstance(profile_keys, list) or len(profile_keys) != profile_count: + raise StarVLAError("FAST policy normalization profiles are invalid") + for index in range(profile_count): + for suffix in ("action_q01", "action_q99", "action_mask"): + values = _gguf_field(reader, f"starvla.normalization.profile.{index}.{suffix}") + if not isinstance(values, list) or len(values) != action_dim: + raise StarVLAError(f"FAST normalization profile {index} is incomplete") + + if expected_metadata is not None: + expected_keys = set(expected_metadata) + actual_keys = { + key for key in reader.fields + if key.startswith("starvla.") or key in expected_keys + } + if actual_keys != expected_keys: + raise StarVLAError("FAST policy GGUF metadata set mismatch") + for key, expected in expected_metadata.items(): + if not _metadata_matches(_gguf_field(reader, key), expected): + raise StarVLAError(f"FAST policy GGUF metadata mismatch: {key}") + if expected_arrays is not None: + actual_arrays = { + ACTION_TOKEN_MAP_TENSOR: action_map, + CODEC_TOKEN_OFFSETS_TENSOR: offsets, + CODEC_TOKEN_BYTES_TENSOR: token_bytes, + } + if set(expected_arrays) != set(actual_arrays) or any( + not np.array_equal(actual_arrays[name], expected) + for name, expected in expected_arrays.items() + ): + raise StarVLAError("FAST policy tensors differ from compiled assets") + + tensor_contract = { + name: {"dtype": dtype, "shape": shape} + for name, (dtype, shape) in expected_shapes.items() + } + record = { + "path": path.name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + "dtype": "integer_runtime_constants", + "architecture": "starvla-policy", + "tensor_count": len(tensors), + "tensor_dtypes": {"I32": 2, "I8": 1}, + "tensor_contract": tensor_contract, + } + del reader + return record + + +def build_bundle_manifest( + *, + manifest: Mapping[str, Any], + entry: Mapping[str, Any], + codec: Mapping[str, Any], + text_component: Mapping[str, Any], + mmproj_component: Mapping[str, Any], + policy_component: Mapping[str, Any], +) -> dict[str, Any]: + if policy_component.get("path") != POLICY_FILENAME: + raise StarVLAError("FAST policy component has an unexpected filename") + return { + "schema_version": 1, + "kind": "starvla_qwen25_fast_gguf_bundle", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": manifest["bundle_uuid"], + "source": portable_source_record(manifest["source"], entry), + "generation": dict(GENERATION_CONTRACT), + "action_token_mapping": manifest["action_token_mapping"], + "fast_codec": { + **dict(codec), + "runtime_storage": "embedded_integer_tensors_in_policy_gguf", + "runtime_policy_gguf": POLICY_FILENAME, + "external_sidecars_required": False, + }, + "components": { + "text": dict(text_component), + "mmproj": dict(mmproj_component), + "policy": dict(policy_component), + }, + "policy_implementation": "finetuned_autoregressive_qwen2_5_vl", + "separate_policy_gguf": POLICY_FILENAME, + } + + +def validate_checkpoint_inventory(records: Sequence[Any]) -> dict[str, Any]: + summary = inventory_summary(list(records)) + mismatches = [ + f"{key}: expected {value!r}, got {summary.get(key)!r}" + for key, value in EXPECTED_INVENTORY.items() + if summary.get(key) != value + ] + by_name = {record.destination_name: record for record in records} + required_shapes = { + "model.embed_tokens.weight": [153713, 2048], + "lm_head.weight": [153713, 2048], + "visual.patch_embed.proj.weight": [1280, 3, 2, 14, 14], + } + for name, shape in required_shapes.items(): + record = by_name.get(name) + if record is None: + mismatches.append(f"{name}: missing") + elif record.shape != shape: + mismatches.append(f"{name}: expected {shape}, got {record.shape}") + if mismatches: + raise StarVLAError( + "Qwen2.5 FAST checkpoint inventory mismatch: " + "; ".join(mismatches) + ) + return summary + + +def preflight( + catalog: Mapping[str, Any], + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, +) -> dict[str, Any]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + policy_hashes = verify_catalog_files(source_dir, entry) + qwen_hashes = verify_catalog_files(qwen_dir, qwen_entry) + qwen_config = validate_qwen_config(qwen_dir) + qwen_processor = validate_qwen_processor(qwen_dir) + generation = validate_generation_config(qwen_dir) + mapping = validate_action_token_mapping(qwen_dir) + codec = validate_fast_codec(codec_dir, codec_entry) + runtime_arrays = compile_fast_runtime_tensors(qwen_dir, codec_dir) + codec["runtime_tensors"] = { + name: { + "dtype": "I32" if array.dtype == np.int32 else "I8", + "shape": list(array.shape), + "sha256": _sha256_bytes(array.tobytes(order="C")), + } + for name, array in runtime_arrays.items() + } + return { + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "source": { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "metadata": policy_hashes, + }, + "qwen": { + "repo_id": qwen_entry["repo_id"], + "revision": qwen_entry["revision"], + "metadata": qwen_hashes, + "config": qwen_config, + "processor": qwen_processor, + }, + "generation": generation, + "action_token_mapping": mapping, + "fast_codec": codec, + } + + +def effective_fast_config(source_dir: Path) -> dict[str, Any]: + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to resolve the FAST config") from exc + try: + source = yaml.safe_load((source_dir / "config.yaml").read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load FAST config.yaml: {exc}") from exc + if not isinstance(source, dict): + raise StarVLAError("FAST config.yaml must contain an object") + framework = source.get("framework") + datasets = source.get("datasets") + if not isinstance(framework, dict) or not isinstance(datasets, dict): + raise StarVLAError("FAST config.yaml is missing framework or datasets") + action = framework.get("action_model") + vla = datasets.get("vla_data") + if not isinstance(action, dict) or not isinstance(vla, dict): + raise StarVLAError("FAST config.yaml is missing action_model or vla_data") + action_dim = action.get("action_dim") + future_window = action.get("future_action_window_size") + cot_prompt = vla.get("CoT_prompt") + image_names = vla.get("obs") + image_size = vla.get("image_size") + if type(action_dim) is not int or type(future_window) is not int: + raise StarVLAError("FAST action dimensions must be integers") + if not isinstance(cot_prompt, str) or not cot_prompt: + raise StarVLAError("FAST CoT_prompt must be a non-empty string") + if not isinstance(image_names, list) or not image_names or any( + not isinstance(name, str) or not name for name in image_names + ): + raise StarVLAError("FAST obs must be a non-empty list of image names") + if ( + not isinstance(image_size, list) + or len(image_size) != 2 + or any(type(value) is not int or value <= 0 for value in image_size) + ): + raise StarVLAError("FAST image_size must contain two positive integers") + return { + "schema_version": 1, + "framework": str(framework.get("framework_py", "QwenFast")), + "backbone": BACKBONE, + "action_model": "autoregressive_vlm_lm_head", + "action_dim": action_dim, + "action_horizon": future_window + 1, + "cot_prompt": cot_prompt, + "image_count": len(image_names), + "image_names": image_names, + "image_size": image_size, + "generation": dict(GENERATION_CONTRACT), + "source_config_sha256": sha256_file(source_dir / "config.yaml"), + } + + +def stage_checkpoint( + *, + checkpoint: Path, + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, + staging_dir: Path, + catalog: Mapping[str, Any], + max_shard_size: int, +) -> dict[str, Any]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + report = preflight(catalog, source_dir, qwen_dir, codec_dir) + verify_checkpoint_file(checkpoint, entry) + if staging_dir.exists(): + raise StarVLAError(f"refusing to overwrite staging directory: {staging_dir}") + staging_dir.parent.mkdir(parents=True, exist_ok=True) + staging_dir.mkdir() + try: + state_dict = load_checkpoint_state(checkpoint) + records = build_inventory(state_dict, entry, enforce_expected=False) + inventory = validate_checkpoint_inventory(records) + validate_qwen_vlm_destination_names( + qwen_dir, + qwen_entry, + records, + backbone=BACKBONE, + ) + + hf_dir = staging_dir / "hf" + source_assets_dir = staging_dir / "source" + hf_dir.mkdir() + source_assets_dir.mkdir() + qwen_assets = copy_qwen_assets(qwen_dir, hf_dir, qwen_entry) + expected_qwen_assets = staged_qwen_asset_hashes(qwen_entry) + if qwen_assets != expected_qwen_assets: + raise StarVLAError("staged Qwen asset hashes do not match the catalog overrides") + policy_assets = copy_policy_assets(source_dir, source_assets_dir, entry) + + effective_path = source_assets_dir / "effective_config.json" + atomic_write_json(effective_path, effective_fast_config(source_dir)) + vlm_output = write_safetensor_shards( + hf_dir, + "model", + "model.safetensors.index.json", + list(records), + state_dict, + max_shard_size, + ) + del state_dict + + codec = validate_fast_codec(codec_dir, codec_entry) + manifest = { + "schema_version": 1, + "kind": "starvla_qwen25_fast_checkpoint_staging", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": bundle_uuid(entry, catalog), + "source": { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "checkpoint": str(checkpoint.resolve()), + "checkpoint_size": checkpoint.stat().st_size, + "checkpoint_sha256": entry["checkpoint"]["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + "qwen_asset": QWEN_ASSET_KEY, + }, + "inventory": inventory, + "qwen_assets": qwen_assets, + "policy_assets": policy_assets, + "fast_codec": codec, + "generation": dict(GENERATION_CONTRACT), + "action_token_mapping": report["action_token_mapping"], + "effective_config": { + "path": "source/effective_config.json", + "size": effective_path.stat().st_size, + "sha256": sha256_file(effective_path), + }, + "vlm_output": vlm_output, + "tensors": [record.to_json() for record in records], + } + atomic_write_json( + staging_dir / STAGING_MANIFEST_FILENAME, + manifest, + overwrite=False, + ) + return manifest + except BaseException: + shutil.rmtree(staging_dir, ignore_errors=True) + raise + + +def validate_staging_manifest( + manifest: Mapping[str, Any], + catalog: Mapping[str, Any], + staging_dir: Path, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + entry, qwen_entry, codec_entry = validate_catalog_contract(catalog) + expected = { + "schema_version": 1, + "kind": "starvla_qwen25_fast_checkpoint_staging", + "variant": VARIANT_KEY, + "framework": FRAMEWORK, + "backbone": BACKBONE, + "model_type": MODEL_TYPE, + "bundle_uuid": bundle_uuid(entry, catalog), + } + mismatches = [ + f"{key}: expected {value!r}, got {manifest.get(key)!r}" + for key, value in expected.items() + if manifest.get(key) != value + ] + source = manifest.get("source") + if not isinstance(source, Mapping): + mismatches.append("source: missing") + else: + source_expected = { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "checkpoint_size": entry["checkpoint"]["size"], + "checkpoint_sha256": entry["checkpoint"]["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + "qwen_asset": QWEN_ASSET_KEY, + } + mismatches.extend( + f"source.{key}: expected {value!r}, got {source.get(key)!r}" + for key, value in source_expected.items() + if source.get(key) != value + ) + inventory = manifest.get("inventory") + if not isinstance(inventory, Mapping): + mismatches.append("inventory: missing") + else: + mismatches.extend( + f"inventory.{key}: expected {value!r}, got {inventory.get(key)!r}" + for key, value in EXPECTED_INVENTORY.items() + if inventory.get(key) != value + ) + if manifest.get("qwen_assets") != staged_qwen_asset_hashes(qwen_entry): + mismatches.append("qwen_assets: mismatch") + expected_policy_assets = { + relative: record["sha256"] for relative, record in entry["file_hashes"].items() + } + if manifest.get("policy_assets") != expected_policy_assets: + mismatches.append("policy_assets: mismatch") + if manifest.get("generation") != GENERATION_CONTRACT: + mismatches.append("generation: mismatch") + expected_mapping = { + "count": ACTION_TOKEN_COUNT, + "fast_token_min": 0, + "fast_token_max": ACTION_TOKEN_COUNT - 1, + "vlm_token_min": ACTION_TOKEN_MIN, + "vlm_token_max": ACTION_TOKEN_MAX, + "mapping": "vlm_token_id = fast_token_id + 151665", + "sha256": qwen_entry["file_hashes"]["added_token_id_map.json"]["sha256"], + } + if manifest.get("action_token_mapping") != expected_mapping: + mismatches.append("action_token_mapping: mismatch") + expected_codec = { + "repo_id": codec_entry["repo_id"], + "revision": codec_entry["revision"], + "scale": 10, + "min_token": -354, + "vocab_size": ACTION_TOKEN_COUNT, + "action_dim": ACTION_DIM, + "time_horizon": ACTION_HORIZON, + "files": { + relative: record["sha256"] + for relative, record in codec_entry["file_hashes"].items() + }, + } + if manifest.get("fast_codec") != expected_codec: + mismatches.append("fast_codec: mismatch") + tensors = manifest.get("tensors") + if not isinstance(tensors, list) or len(tensors) != EXPECTED_INVENTORY["total_tensors"]: + mismatches.append("tensors: incomplete checkpoint inventory") + if mismatches: + raise StarVLAError("invalid Qwen2.5 FAST staging manifest: " + "; ".join(mismatches)) + + hf_dir = staging_dir / "hf" + verify_staged_assets(hf_dir, manifest["qwen_assets"], component="Qwen") + verify_staged_assets( + staging_dir / "source", + manifest["policy_assets"], + component="FAST source", + ) + effective = manifest.get("effective_config") + if not isinstance(effective, Mapping): + raise StarVLAError("FAST staging manifest has no effective_config record") + if effective.get("path") != "source/effective_config.json": + raise StarVLAError("FAST staging manifest has an invalid effective_config path") + effective_path = staging_dir / "source" / "effective_config.json" + if ( + not effective_path.is_file() + or effective_path.stat().st_size != effective.get("size") + or sha256_file(effective_path) != effective.get("sha256") + ): + raise StarVLAError("FAST staged effective_config size/SHA256 mismatch") + index = verify_staged_shards(hf_dir, manifest["vlm_output"], component="FAST VLM") + staged_names = set(index["weight_map"]) + manifest_names = { + str(record.get("destination_name")) + for record in tensors + if isinstance(record, Mapping) + } + if len(manifest_names) != len(tensors) or staged_names != manifest_names: + raise StarVLAError("FAST staged tensor names do not match the checkpoint inventory") + return entry, qwen_entry, codec_entry + + +def _reserve_output_directory(output_dir: Path) -> None: + if output_dir.exists(): + raise StarVLAError(f"refusing to overwrite output directory: {output_dir}") + output_dir.parent.mkdir(parents=True, exist_ok=True) + output_dir.mkdir() + + +def convert_staging( + *, + staging_dir: Path, + source_dir: Path, + qwen_dir: Path, + codec_dir: Path, + output_dir: Path, + catalog: Mapping[str, Any], + llama_root: Path, + python: str, +) -> dict[str, Any]: + manifest_path = staging_dir / STAGING_MANIFEST_FILENAME + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load staging manifest {manifest_path}: {exc}") from exc + entry, _, codec_entry = validate_staging_manifest(manifest, catalog, staging_dir) + preflight(catalog, source_dir, qwen_dir, codec_dir) + verified_llama = verify_llama_checkout( + llama_root.resolve(strict=True), + str(manifest["source"]["llama_cpp_revision"]), + ) + + _reserve_output_directory(output_dir) + try: + text_output = output_dir / TEXT_FILENAME + mmproj_output = output_dir / MMPROJ_FILENAME + text_metadata = output_dir / "text-metadata.json" + mmproj_metadata = output_dir / "mmproj-metadata.json" + common_metadata = { + "general.source.uuid": manifest["bundle_uuid"], + "general.source.url": ( + f"https://huggingface.co/{entry['repo_id']}/tree/{entry['revision']}" + ), + "general.finetune": "starvla-qwen25-fast", + } + atomic_write_json( + text_metadata, + { + **common_metadata, + "general.name": "StarVLA Qwen2.5-VL FAST text policy", + }, + ) + atomic_write_json( + mmproj_metadata, + { + **common_metadata, + "general.name": "StarVLA Qwen2.5-VL FAST mmproj", + }, + ) + commands = build_commands( + python, + staging_dir / "hf", + text_output, + mmproj_output, + text_metadata, + mmproj_metadata, + "bf16", + "bf16", + llama_root=verified_llama, + ) + for command in commands: + subprocess.run(command, check=True, cwd=REPOSITORY_ROOT) + for output in (text_output, mmproj_output): + if not output.is_file() or output.stat().st_size == 0: + raise StarVLAError(f"converter did not create {output}") + text_metadata.unlink() + mmproj_metadata.unlink() + + policy_metadata, policy_arrays = build_fast_runtime_policy( + manifest=manifest, + entry=entry, + codec_entry=codec_entry, + source_dir=source_dir, + qwen_dir=qwen_dir, + codec_dir=codec_dir, + ) + policy_output = output_dir / POLICY_FILENAME + write_fast_runtime_policy_gguf( + policy_output, + policy_metadata, + policy_arrays, + ) + policy_component = validate_fast_runtime_policy_gguf( + policy_output, + expected_metadata=policy_metadata, + expected_arrays=policy_arrays, + ) + bundle = build_bundle_manifest( + manifest=manifest, + entry=entry, + codec=validate_fast_codec(codec_dir, codec_entry), + text_component={ + "path": TEXT_FILENAME, + "size": text_output.stat().st_size, + "sha256": sha256_file(text_output), + "dtype": "bf16", + }, + mmproj_component={ + "path": MMPROJ_FILENAME, + "size": mmproj_output.stat().st_size, + "sha256": sha256_file(mmproj_output), + "dtype": "bf16", + }, + policy_component=policy_component, + ) + atomic_write_json( + output_dir / BUNDLE_MANIFEST_FILENAME, + bundle, + overwrite=False, + ) + return bundle + except BaseException: + shutil.rmtree(output_dir, ignore_errors=True) + raise + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--qwen-assets", type=Path, required=True) + parser.add_argument("--fast-codec", type=Path, required=True) + parser.add_argument("--staging-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--llama-root", + type=Path, + default=REPOSITORY_ROOT / "third_party" / "llama.cpp", + ) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--max-shard-size", type=parse_size, default=parse_size("2G")) + parser.add_argument("--preflight", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--stage-only", action="store_true") + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + try: + catalog = load_catalog(args.catalog) + report = preflight( + catalog, + args.source_dir, + args.qwen_assets, + args.fast_codec, + ) + if args.preflight: + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + if args.dry_run: + output_dir = args.output_dir or Path("ckpts/starvla/gguf/qwen25_fast") + commands = build_commands( + args.python, + args.staging_dir / "hf", + output_dir / TEXT_FILENAME, + output_dir / MMPROJ_FILENAME, + output_dir / "text-metadata.json", + output_dir / "mmproj-metadata.json", + "bf16", + "bf16", + llama_root=args.llama_root, + ) + print( + json.dumps( + { + **report, + "checkpoint_required_for_execution": True, + "commands": commands, + "bundle_components": { + "text": TEXT_FILENAME, + "mmproj": MMPROJ_FILENAME, + "policy": POLICY_FILENAME, + }, + "runtime_policy": { + "built_in_process": True, + "external_sidecars_required": False, + "tensor_count": len(FAST_RUNTIME_TENSOR_NAMES), + }, + }, + indent=2, + sort_keys=True, + ) + ) + return 0 + if args.checkpoint is None: + raise StarVLAError("--checkpoint is required unless --preflight or --dry-run is used") + manifest = stage_checkpoint( + checkpoint=args.checkpoint, + source_dir=args.source_dir, + qwen_dir=args.qwen_assets, + codec_dir=args.fast_codec, + staging_dir=args.staging_dir, + catalog=catalog, + max_shard_size=args.max_shard_size, + ) + print(f"staging manifest: {args.staging_dir / STAGING_MANIFEST_FILENAME}") + if args.stage_only: + print(json.dumps(manifest["inventory"], indent=2, sort_keys=True)) + return 0 + if args.output_dir is None: + raise StarVLAError("--output-dir is required unless --stage-only is used") + bundle = convert_staging( + staging_dir=args.staging_dir, + source_dir=args.source_dir, + qwen_dir=args.qwen_assets, + codec_dir=args.fast_codec, + output_dir=args.output_dir, + catalog=catalog, + llama_root=args.llama_root, + python=args.python, + ) + print(f"bundle manifest: {args.output_dir / BUNDLE_MANIFEST_FILENAME}") + print(json.dumps(bundle["components"], indent=2, sort_keys=True)) + return 0 + except ( + StarVLAError, + OSError, + json.JSONDecodeError, + subprocess.CalledProcessError, + KeyError, + ) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py new file mode 100755 index 0000000..4243dcb --- /dev/null +++ b/tools/hf2gguf/starvla/convert_starvla_qwen_to_gguf.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Invoke the pinned llama.cpp converter for StarVLA Qwen-VL text and mmproj.""" + +from __future__ import annotations + +import argparse +import json +import re +import subprocess +import sys +from pathlib import Path + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + atomic_write_json, + default_mmproj_filename, + default_text_filename, + get_variant, + load_catalog, + validate_surgery_manifest, + verify_staged_assets, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_ROOT = REPOSITORY_ROOT / "third_party" / "llama.cpp" +LLAMA_CONVERTER = LLAMA_ROOT / "convert_hf_to_gguf.py" +LLAMA_GGUF_PY = LLAMA_ROOT / "gguf-py" +PINNED_REVISION_RE = re.compile(r"[0-9a-f]{40}") + + +def git_revision(path: Path) -> str: + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to resolve git revision for {path}: {exc}") from exc + return result.stdout.strip() + + +def git_worktree_changes(path: Path) -> str: + try: + result = subprocess.run( + [ + "git", + "-C", + str(path), + "status", + "--porcelain=v1", + "--untracked-files=all", + ], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to inspect git worktree for {path}: {exc}") from exc + return result.stdout.strip() + + +def canonical_llama_root(path: Path) -> Path: + """Require an explicit, canonical llama.cpp checkout root with converter sources.""" + if not path.is_absolute(): + raise StarVLAError(f"llama.cpp root must be an absolute canonical directory: {path}") + try: + canonical = path.resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise StarVLAError(f"failed to resolve llama.cpp root {path}: {exc}") from exc + if canonical != path or not canonical.is_dir(): + raise StarVLAError(f"llama.cpp root must be an absolute canonical directory: {path}") + + converter = canonical / "convert_hf_to_gguf.py" + gguf_py = canonical / "gguf-py" + if not converter.is_file(): + raise StarVLAError(f"missing llama.cpp converter: {converter}") + if not gguf_py.is_dir(): + raise StarVLAError(f"missing llama.cpp gguf-py directory: {gguf_py}") + return canonical + + +def verify_llama_checkout(path: Path, expected_revision: str) -> Path: + """Verify that path is the clean root of the exact manifest-pinned checkout.""" + root = canonical_llama_root(path) + if PINNED_REVISION_RE.fullmatch(expected_revision) is None: + raise StarVLAError( + f"manifest contains an invalid pinned llama.cpp revision: {expected_revision!r}" + ) + + try: + result = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--show-toplevel"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError) as exc: + raise StarVLAError(f"failed to resolve git root for {root}: {exc}") from exc + try: + git_root = Path(result.stdout.strip()).resolve(strict=True) + except (OSError, RuntimeError) as exc: + raise StarVLAError(f"failed to resolve git root reported for {root}: {exc}") from exc + if git_root != root: + raise StarVLAError( + f"llama.cpp root must be the canonical Git worktree root: expected {root}, got {git_root}" + ) + + actual_revision = git_revision(root) + if actual_revision != expected_revision: + raise StarVLAError( + f"llama.cpp revision mismatch: expected {expected_revision}, got {actual_revision}; " + "use the revision pinned by the checkpoint catalog" + ) + worktree_changes = git_worktree_changes(root) + if worktree_changes: + raise StarVLAError( + "llama.cpp has tracked or untracked worktree changes; " + f"use the clean pinned revision for conversion:\n{worktree_changes}" + ) + return root + + +def build_commands( + python: str, + hf_dir: Path, + text_output: Path, + mmproj_output: Path, + text_metadata: Path, + mmproj_metadata: Path, + text_dtype: str, + mmproj_dtype: str, + *, + llama_root: Path = LLAMA_ROOT, +) -> list[list[str]]: + # Isolated mode excludes the working directory, PYTHONPATH and user site + # from imports while the pinned converter adds its own gguf-py directory. + converter = llama_root / "convert_hf_to_gguf.py" + common = [python, "-I", str(converter), str(hf_dir)] + return [ + common + + [ + "--outfile", + str(text_output), + "--outtype", + text_dtype, + "--metadata", + str(text_metadata), + ], + common + + [ + "--outfile", + str(mmproj_output), + "--outtype", + mmproj_dtype, + "--metadata", + str(mmproj_metadata), + "--mmproj", + ], + ] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--llama-root", type=Path, required=True) + parser.add_argument("--text-filename") + parser.add_argument("--mmproj-filename") + parser.add_argument( + "--text-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_TEXT_DTYPE, + ) + parser.add_argument( + "--mmproj-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_MMPROJ_DTYPE, + ) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if not args.hf_dir.is_dir(): + raise StarVLAError(f"missing HF staging directory: {args.hf_dir}") + try: + manifest = json.loads(args.surgery_manifest.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load surgery manifest {args.surgery_manifest}: {exc}") from exc + + catalog = load_catalog(args.catalog) + variant_name = str(manifest.get("variant", "")) + variant = get_variant(catalog, variant_name) + validate_surgery_manifest(manifest, variant, catalog) + verify_staged_assets(args.hf_dir, manifest.get("qwen_assets", {}), component="Qwen") + expected_revision = str(manifest.get("source", {}).get("llama_cpp_revision", "")) + llama_root = verify_llama_checkout(args.llama_root, expected_revision) + + args.output_dir.mkdir(parents=True, exist_ok=True) + text_filename = args.text_filename or default_text_filename(variant_name, args.text_dtype) + mmproj_filename = args.mmproj_filename or default_mmproj_filename( + variant_name, args.mmproj_dtype + ) + text_output = args.output_dir / text_filename + mmproj_output = args.output_dir / mmproj_filename + text_metadata = args.output_dir / "text-metadata.json" + mmproj_metadata = args.output_dir / "mmproj-metadata.json" + bundle_uuid = str(manifest["bundle_uuid"]) + source = manifest["source"] + backbone = str(manifest.get("backbone", variant["backbone"])) + backbone_label = { + "qwen3_vl": "Qwen3-VL", + "qwen2_5_vl": "Qwen2.5-VL", + }.get(backbone) + if backbone_label is None: + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + common_metadata = { + "general.source.uuid": bundle_uuid, + "general.source.url": f"https://huggingface.co/{source['repo_id']}/tree/{source['revision']}", + "general.finetune": f"starvla-{manifest['variant']}", + } + atomic_write_json( + text_metadata, + { + **common_metadata, + "general.name": f"StarVLA {backbone_label} {manifest['variant']} text", + }, + ) + atomic_write_json( + mmproj_metadata, + { + **common_metadata, + "general.name": f"StarVLA {backbone_label} {manifest['variant']} mmproj", + }, + ) + + commands = build_commands( + args.python, + args.hf_dir, + text_output, + mmproj_output, + text_metadata, + mmproj_metadata, + args.text_dtype, + args.mmproj_dtype, + llama_root=llama_root, + ) + if args.dry_run: + print(json.dumps(commands, indent=2)) + return 0 + + for command in commands: + subprocess.run(command, check=True, cwd=REPOSITORY_ROOT) + for output in (text_output, mmproj_output): + if not output.is_file() or output.stat().st_size == 0: + raise StarVLAError(f"llama.cpp converter did not create the expected output: {output}") + print(f"text GGUF: {text_output}") + print(f"mmproj GGUF: {mmproj_output}") + return 0 + except (StarVLAError, OSError, subprocess.CalledProcessError, KeyError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/download_starvla.py b/tools/hf2gguf/starvla/download_starvla.py new file mode 100755 index 0000000..364f2b5 --- /dev/null +++ b/tools/hf2gguf/starvla/download_starvla.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Download pinned StarVLA sources and shared tokenizer assets.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Any, Sequence + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + atomic_write_json, + get_variant, + load_catalog, + sha256_file, +) + + +DEFAULT_BACKBONE = "qwen3_vl" + + +def destination_for(root: Path, entry: dict[str, Any]) -> Path: + return root / str(entry["directory"]) / str(entry["revision"]) + + +def variant_backbone(entry: dict[str, Any]) -> str: + return str(entry["backbone"]) + + +def available_backbones(catalog: dict[str, Any]) -> list[str]: + return list( + dict.fromkeys( + variant_backbone(entry) + for entry in catalog["variants"].values() + ) + ) + + +def resolve_variant_keys( + catalog: dict[str, Any], + requested: Sequence[str] | None, + requested_backbone: str | None, +) -> tuple[str, list[str]]: + variants = catalog["variants"] + backbones = available_backbones(catalog) + if requested_backbone is not None: + backbone = requested_backbone + else: + direct_backbones = { + variant_backbone(variants[name]) + for name in (requested or ()) + if name in variants + } + if len(direct_backbones) > 1: + raise StarVLAError( + "requested variants span multiple backbones; select one with --backbone" + ) + backbone = next(iter(direct_backbones), DEFAULT_BACKBONE) + if backbone not in backbones: + raise StarVLAError( + f"unknown StarVLA backbone {backbone!r}; expected one of {backbones}" + ) + + candidates = { + name: entry + for name, entry in variants.items() + if variant_backbone(entry) == backbone + } + tokens = list(requested or ("oft",)) + if "all" in tokens: + if len(tokens) != 1: + raise StarVLAError( + "--variant all cannot be combined with another variant" + ) + return backbone, list(candidates) + + selected: list[str] = [] + for token in tokens: + if token in candidates: + key = token + else: + matches = [ + name + for name, entry in candidates.items() + if entry.get("framework") == token + ] + if not matches: + accepted = sorted( + { + *candidates, + *(str(entry["framework"]) for entry in candidates.values()), + "all", + } + ) + raise StarVLAError( + f"variant {token!r} is not available for backbone {backbone!r}; " + f"expected one of {accepted}" + ) + if len(matches) != 1: + raise StarVLAError( + f"framework alias {token!r} is ambiguous for backbone {backbone!r}; " + f"use one of {matches}" + ) + key = matches[0] + if key not in selected: + selected.append(key) + return backbone, selected + + +def required_shared_assets( + catalog: dict[str, Any], + variant_keys: Sequence[str], +) -> list[str]: + names: list[str] = [] + has_fast = False + for variant_key in variant_keys: + entry = get_variant(catalog, variant_key) + qwen_asset = str(entry["qwen_asset"]) + if qwen_asset not in catalog["shared_assets"]: + raise StarVLAError( + f"variant {variant_key!r} references unknown Qwen asset {qwen_asset!r}" + ) + if qwen_asset not in names: + names.append(qwen_asset) + has_fast = has_fast or entry.get("framework") == "fast" + if has_fast and "fast_codec" not in names: + names.append("fast_codec") + return names + + +def download_entry( + entry: dict[str, Any], + root: Path, + files: list[str], + *, + dry_run: bool, + local_files_only: bool, + force_download: bool, +) -> dict[str, Any]: + destination = destination_for(root, entry) + result = { + "repo_id": entry["repo_id"], + "revision": entry["revision"], + "directory": (Path(str(entry["directory"])) / str(entry["revision"])).as_posix(), + "requested_files": files, + "files": [], + } + if dry_run: + return result + + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise StarVLAError("huggingface_hub is required; install tools/hf2gguf/environment.yaml") from exc + + destination.mkdir(parents=True, exist_ok=True) + try: + snapshot_download( + repo_id=str(entry["repo_id"]), + revision=str(entry["revision"]), + allow_patterns=files, + local_dir=destination, + local_files_only=local_files_only, + force_download=force_download, + ) + except Exception as exc: + raise StarVLAError(f"failed to download {entry['repo_id']}@{entry['revision']}: {exc}") from exc + + missing = [relative for relative in files if not (destination / relative).is_file()] + if missing: + raise StarVLAError(f"download completed with missing files in {destination}: {missing}") + + expected_records = dict(entry.get("file_hashes", {})) + expected_records.update(entry.get("optional_weight_hashes", {})) + checkpoint = entry.get("checkpoint") + if checkpoint is not None: + expected_records[str(checkpoint["path"])] = checkpoint + for relative in sorted(files): + path = destination / relative + expected = expected_records.get(relative) + if expected is None: + raise StarVLAError(f"catalog has no pinned size/SHA256 for requested file: {relative}") + record = {"path": relative, "size": path.stat().st_size, "sha256": sha256_file(path)} + if record["size"] != expected["size"] or record["sha256"] != expected["sha256"]: + raise StarVLAError( + f"downloaded file size/SHA256 mismatch for {path}: " + f"expected {expected['size']}/{expected['sha256']}, " + f"got {record['size']}/{record['sha256']}" + ) + result["files"].append(record) + return result + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + action="append", + help=( + "catalog variant key or framework alias to download " + "(repeatable; 'all' selects every variant for the backbone; " + "default: OFT for the selected backbone)" + ), + ) + parser.add_argument( + "--backbone", + help=( + "backbone selector from the catalog " + f"(default: infer from exact variant keys, otherwise {DEFAULT_BACKBONE})" + ), + ) + parser.add_argument("--root", type=Path, default=Path("ckpts/starvla/sources")) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--metadata-only", + action="store_true", + help="skip policy checkpoints and all optional base weights", + ) + parser.add_argument( + "--skip-checkpoint", + action="store_true", + help="skip policy checkpoints while retaining requested optional base weights", + ) + parser.add_argument( + "--include-base-weights", + action="store_true", + help="download optional safetensor shards for the selected Qwen base assets", + ) + parser.add_argument( + "--include-fast-weights", + action="store_true", + help=( + "download the action-ready Qwen weights for a selected FAST variant; " + "the policy checkpoint is downloaded separately" + ), + ) + parser.add_argument("--no-shared-assets", action="store_true") + parser.add_argument("--local-files-only", action="store_true") + parser.add_argument("--force-download", action="store_true") + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args(argv) + + +def main() -> int: + args = parse_args() + try: + catalog = load_catalog(args.catalog) + backbone, variants = resolve_variant_keys(catalog, args.variant, args.backbone) + + manifest: dict[str, Any] = { + "schema_version": 1, + "catalog_sha256": sha256_file(args.catalog), + "source_revisions": catalog["source_revisions"], + "backbone": backbone, + "variants": variants, + "metadata_only": bool(args.metadata_only), + "skip_checkpoint": bool(args.skip_checkpoint), + "downloads": {}, + } + + if not args.no_shared_assets: + shared_names = required_shared_assets(catalog, variants) + variant_entries = { + variant: get_variant(catalog, variant) + for variant in variants + } + fast_qwen_assets = { + str(entry["qwen_asset"]) + for entry in variant_entries.values() + if entry.get("framework") == "fast" + } + for name in shared_names: + raw_entry = catalog["shared_assets"][name] + entry = dict(raw_entry) + files = list(entry["files"]) + include_optional_weights = ( + args.include_base_weights + or (args.include_fast_weights and name in fast_qwen_assets) + ) + if include_optional_weights and not args.metadata_only: + files.extend(entry.get("optional_weight_files", [])) + manifest["downloads"][f"asset:{name}"] = download_entry( + entry, + args.root, + files, + dry_run=args.dry_run, + local_files_only=args.local_files_only, + force_download=args.force_download, + ) + + for variant in variants: + entry = get_variant(catalog, variant) + files = list(entry.get("files", [])) + checkpoint = entry["checkpoint"] + if not args.metadata_only and not args.skip_checkpoint: + files.append(str(checkpoint["path"])) + download = download_entry( + entry, + args.root, + files, + dry_run=args.dry_run, + local_files_only=args.local_files_only, + force_download=args.force_download, + ) + manifest["downloads"][f"variant:{variant}"] = download + + if not args.metadata_only and not args.skip_checkpoint and not args.dry_run: + record = next( + (item for item in download["files"] if item["path"] == checkpoint["path"]), + None, + ) + if record is None: + raise StarVLAError(f"download manifest has no checkpoint record for {checkpoint['path']}") + if record["size"] != checkpoint["size"] or record["sha256"] != checkpoint["sha256"]: + raise StarVLAError( + f"checkpoint verification failed for {checkpoint['path']}: " + f"expected size/hash {checkpoint['size']}/{checkpoint['sha256']}, " + f"got {record['size']}/{record['sha256']}" + ) + + if args.dry_run: + import json + + print(json.dumps(manifest, indent=2, sort_keys=True)) + else: + manifest_path = args.root / "download_manifest.json" + atomic_write_json(manifest_path, manifest) + print(f"download manifest: {manifest_path}") + return 0 + except StarVLAError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/environment.yaml b/tools/hf2gguf/starvla/environment.yaml new file mode 100644 index 0000000..5052ab7 --- /dev/null +++ b/tools/hf2gguf/starvla/environment.yaml @@ -0,0 +1,23 @@ +name: starvla_gguf_converter +channels: + - conda-forge +dependencies: + - python=3.11 + - pip + - pip: + - torch==2.6.0 + - torchvision==0.21.0 + - numpy==1.26.4 + - safetensors==0.7.0 + - sentencepiece + - transformers==4.57.0 + - tokenizers==0.22.2 + - accelerate==1.5.2 + - diffusers==0.37.1 + - omegaconf==2.3.0 + - pillow==12.1.1 + - qwen-vl-utils==0.0.14 + - rich + - scipy + - huggingface_hub>=0.36.0 + - pyyaml diff --git a/tools/hf2gguf/starvla/starvla_checkpoint.py b/tools/hf2gguf/starvla/starvla_checkpoint.py new file mode 100755 index 0000000..a9722ee --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_checkpoint.py @@ -0,0 +1,1101 @@ +#!/usr/bin/env python3 +"""Shared catalog and strict checkpoint inventory helpers for StarVLA.""" + +from __future__ import annotations + +import hashlib +import json +import os +import uuid +from copy import deepcopy +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path, PureWindowsPath +from typing import Any, Mapping, Sequence + + +DEFAULT_CATALOG = Path(__file__).with_name("checkpoint_catalog.json") +DEFAULT_TEXT_DTYPE = "bf16" +DEFAULT_MMPROJ_DTYPE = "bf16" +DEFAULT_POLICY_DTYPE = "fp32" + +STARVLA_ARTIFACT_STEMS = { + "oft": "oft", + "groot": "groot", + "pi_v3": "pi-v3", + "qwen25_oft": "qwen25-oft", + "qwen25_groot": "qwen25-groot", + "qwen25_pi": "qwen25-pi", + "qwen25_fast": "qwen25-fast", +} + +LEGACY_QWEN3_ASSET = "qwen3_vl_4b_instruct" +SUPPORTED_BACKBONES = {"qwen3_vl", "qwen2_5_vl"} +SUPPORTED_FRAMEWORKS = {"oft", "groot", "pi", "pi_v3", "fast"} +GENERATED_QWEN_ASSET_PATHS = {"model.safetensors.index.json"} + +VLM_SOURCE_RULES = ( + ("qwen_vl_interface.model.model.visual.", "visual"), + ("qwen_vl_interface.model.model.language_model.", "text"), + ("qwen_vl_interface.model.lm_head.", "lm_head"), +) + +VLM_DESTINATION_PREFIXES = { + "qwen3_vl": { + "visual": "model.visual.", + "text": "model.language_model.", + "lm_head": "lm_head.", + }, + "qwen2_5_vl": { + "visual": "visual.", + "text": "model.", + "lm_head": "lm_head.", + }, +} + + +class StarVLAError(RuntimeError): + """Raised when a catalog or checkpoint violates the conversion contract.""" + + +def artifact_stem(variant: str) -> str: + try: + return STARVLA_ARTIFACT_STEMS[variant] + except KeyError as exc: + raise StarVLAError(f"unsupported StarVLA artifact variant: {variant!r}") from exc + + +def default_text_filename(variant: str, dtype: str = DEFAULT_TEXT_DTYPE) -> str: + return f"qwen-{artifact_stem(variant)}-{dtype}.gguf" + + +def default_mmproj_filename(variant: str, dtype: str = DEFAULT_MMPROJ_DTYPE) -> str: + return f"mmproj-{artifact_stem(variant)}-{dtype}.gguf" + + +@dataclass(frozen=True) +class TensorRecord: + source_name: str + destination_name: str + component: str + role: str + shape: list[int] + dtype: str + numel: int + nbytes: int + storage_offset: int + storage_alias: str | None = None + + def to_json(self) -> dict[str, Any]: + return asdict(self) + + +def _safe_relative_path(value: Any, *, field: str) -> Path: + if not isinstance(value, str): + raise StarVLAError(f"unsafe relative path in {field}: expected a string, got {value!r}") + if not value or "\x00" in value or "\\" in value: + raise StarVLAError(f"unsafe relative path in {field}: {value!r}") + + path = Path(value) + windows_path = PureWindowsPath(value) + if ( + path.is_absolute() + or windows_path.is_absolute() + or bool(windows_path.drive) + or not path.parts + or any(part in ("", ".", "..") for part in path.parts) + or path.as_posix() != value + ): + raise StarVLAError(f"unsafe relative path in {field}: {value!r}") + return path + + +def _validate_revision(value: Any, *, field: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 40 + or any(ch not in "0123456789abcdef" for ch in value) + ): + raise StarVLAError(f"invalid pinned revision in {field}: expected 40 lowercase hex characters") + return value + + +def _validate_sha256(value: Any, *, field: str) -> str: + if ( + not isinstance(value, str) + or len(value) != 64 + or any(ch not in "0123456789abcdef" for ch in value) + ): + raise StarVLAError(f"invalid SHA256 in {field}") + return value + + +def _validate_positive_size(value: Any, *, field: str) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise StarVLAError(f"invalid size in {field}: expected a positive integer") + return value + + +def load_catalog(path: Path | str = DEFAULT_CATALOG) -> dict[str, Any]: + catalog_path = Path(path) + try: + catalog = json.loads(catalog_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load checkpoint catalog {catalog_path}: {exc}") from exc + + if not isinstance(catalog, dict): + raise StarVLAError("checkpoint catalog root must be an object") + if catalog.get("schema_version") != 1: + raise StarVLAError(f"unsupported checkpoint catalog schema: {catalog.get('schema_version')!r}") + + source_revisions = catalog.get("source_revisions") + if not isinstance(source_revisions, dict): + raise StarVLAError("checkpoint catalog source_revisions must be an object") + for source in ("starvla", "llama_cpp"): + _validate_revision(source_revisions.get(source), field=f"source_revisions.{source}") + for source, revision in source_revisions.items(): + _validate_revision(revision, field=f"source_revisions.{source}") + + shared_assets = catalog.get("shared_assets") + if not isinstance(shared_assets, dict): + raise StarVLAError("checkpoint catalog shared_assets must be an object") + variants = catalog.get("variants") + if not isinstance(variants, dict) or not variants: + raise StarVLAError("checkpoint catalog has no variants") + for name, entry in variants.items(): + if not isinstance(entry, dict): + raise StarVLAError(f"catalog variant {name!r} must be an object") + framework = entry.get("framework") + if framework not in SUPPORTED_FRAMEWORKS: + raise StarVLAError( + f"catalog variant {name!r} has unsupported framework={framework!r}" + ) + if entry.get("model_type") != "starvla": + raise StarVLAError(f"catalog variant {name!r} must use model_type='starvla'") + backbone = entry.get("backbone") + if backbone not in SUPPORTED_BACKBONES: + raise StarVLAError( + f"catalog variant {name!r} has unsupported backbone={backbone!r}" + ) + qwen_asset = entry.get("qwen_asset") + if not isinstance(qwen_asset, str) or qwen_asset not in shared_assets: + raise StarVLAError( + f"catalog variant {name!r} references unknown qwen_asset={qwen_asset!r}" + ) + if not isinstance(entry.get("repo_id"), str) or not entry["repo_id"]: + raise StarVLAError(f"catalog variant {name!r} is missing repo_id/revision") + if not isinstance(entry.get("default_unnorm_key"), str) or not entry["default_unnorm_key"]: + raise StarVLAError(f"catalog variant {name!r} has no default_unnorm_key") + _validate_revision(entry.get("revision"), field=f"variant {name}.revision") + checkpoint = entry.get("checkpoint") + if not isinstance(checkpoint, dict): + raise StarVLAError(f"catalog variant {name!r} checkpoint must be an object") + _safe_relative_path(checkpoint.get("path"), field=f"variant {name}.checkpoint.path") + _validate_positive_size(checkpoint.get("size"), field=f"variant {name}.checkpoint.size") + _validate_sha256(checkpoint.get("sha256"), field=f"variant {name}.checkpoint.sha256") + entries = { + **{f"shared asset {name}": entry for name, entry in shared_assets.items()}, + **{f"variant {name}": entry for name, entry in variants.items()}, + } + for label, entry in entries.items(): + if not isinstance(entry, dict): + raise StarVLAError(f"catalog {label} must be an object") + _safe_relative_path(entry.get("directory"), field=f"{label}.directory") + if not isinstance(entry.get("repo_id"), str) or not entry["repo_id"]: + raise StarVLAError(f"catalog {label} has invalid repo_id") + _validate_revision(entry.get("revision"), field=f"{label}.revision") + + files = entry.get("files", []) + file_hashes = entry.get("file_hashes", {}) + if not isinstance(files, list) or any(not isinstance(relative, str) for relative in files): + raise StarVLAError(f"catalog {label} has invalid or duplicate files") + for index, relative in enumerate(files): + _safe_relative_path(relative, field=f"{label}.files[{index}]") + if len(files) != len(set(files)): + raise StarVLAError(f"catalog {label} has invalid or duplicate files") + if not isinstance(file_hashes, dict): + raise StarVLAError(f"catalog {label} file_hashes must be an object") + if set(file_hashes) != set(files): + raise StarVLAError(f"catalog {label} file_hashes must cover files exactly") + for relative, record in file_hashes.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid file record for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.file_hashes[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.file_hashes[{relative!r}].sha256") + staged_overrides = entry.get("staged_overrides", {}) + if not isinstance(staged_overrides, dict) or not set(staged_overrides).issubset(files): + raise StarVLAError(f"catalog {label} has invalid staged_overrides") + for relative, record in staged_overrides.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid staged override for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.staged_overrides[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.staged_overrides[{relative!r}].sha256") + optional_files = entry.get("optional_weight_files", []) + optional_hashes = entry.get("optional_weight_hashes", {}) + if not isinstance(optional_files, list) or any(not isinstance(relative, str) for relative in optional_files): + raise StarVLAError(f"catalog {label} has invalid optional_weight_files") + for index, relative in enumerate(optional_files): + _safe_relative_path(relative, field=f"{label}.optional_weight_files[{index}]") + if len(optional_files) != len(set(optional_files)): + raise StarVLAError(f"catalog {label} has duplicate optional weights") + if not isinstance(optional_hashes, dict) or set(optional_hashes) != set(optional_files): + raise StarVLAError(f"catalog {label} optional_weight_hashes must cover optional weights exactly") + for relative, record in optional_hashes.items(): + if not isinstance(record, dict): + raise StarVLAError(f"catalog {label} has invalid optional weight record for {relative!r}") + _validate_positive_size(record.get("size"), field=f"{label}.optional_weight_hashes[{relative!r}].size") + _validate_sha256(record.get("sha256"), field=f"{label}.optional_weight_hashes[{relative!r}].sha256") + return catalog + + +def get_variant(catalog: Mapping[str, Any], variant: str) -> dict[str, Any]: + variants = catalog.get("variants", {}) + if variant not in variants: + raise StarVLAError(f"unknown StarVLA variant {variant!r}; expected one of {sorted(variants)}") + entry = dict(variants[variant]) + entry["_catalog_key"] = variant + return entry + + +def local_checkpoint_catalog( + catalog: Mapping[str, Any], + variant_name: str, + checkpoint: Path, + source_dir: Path, + default_unnorm_key: str | None = None, +) -> dict[str, Any]: + """Bind a training checkpoint and its run metadata to a catalog variant.""" + if not checkpoint.is_file(): + raise StarVLAError(f"checkpoint does not exist: {checkpoint}") + required_assets = ("config.yaml", "dataset_statistics.json") + for name in required_assets: + if not (source_dir / name).is_file(): + raise StarVLAError(f"training run is missing {name}: {source_dir / name}") + try: + stats = json.loads((source_dir / "dataset_statistics.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load training dataset statistics: {exc}") from exc + if not isinstance(stats, dict) or not stats: + raise StarVLAError("training dataset_statistics.json must contain at least one profile") + + result = deepcopy(catalog) + variant = get_variant(result, variant_name) + selected_key = default_unnorm_key or str(variant["default_unnorm_key"]) + if selected_key not in stats: + if default_unnorm_key is None and len(stats) == 1: + selected_key = next(iter(stats)) + else: + raise StarVLAError( + f"normalization profile {selected_key!r} is not present; " + f"choose one of {sorted(stats)}" + ) + + local_entry = result["variants"][variant_name] + checkpoint_sha256 = sha256_file(checkpoint) + local_entry["repo_id"] = "local" + local_entry["revision"] = checkpoint_sha256[:40] + local_entry["checkpoint"] = { + "path": checkpoint.name, + "size": checkpoint.stat().st_size, + "sha256": checkpoint_sha256, + } + local_entry["files"] = list(required_assets) + local_entry["file_hashes"] = { + name: { + "size": (source_dir / name).stat().st_size, + "sha256": sha256_file(source_dir / name), + } + for name in required_assets + } + local_entry["default_unnorm_key"] = selected_key + return result + + +def portable_source_record( + source: Mapping[str, Any], variant_entry: Mapping[str, Any] +) -> dict[str, Any]: + """Replace the staging checkpoint path with its catalog-relative path.""" + checkpoint = variant_entry.get("checkpoint") + if not isinstance(checkpoint, Mapping): + raise StarVLAError("catalog variant has no policy checkpoint") + result = dict(source) + result["checkpoint"] = _safe_relative_path( + checkpoint.get("path"), field="variant checkpoint path" + ).as_posix() + return result + + +def get_qwen_asset( + catalog: Mapping[str, Any], variant_entry: Mapping[str, Any] +) -> tuple[str, dict[str, Any]]: + asset_name = variant_entry.get("qwen_asset") + shared_assets = catalog.get("shared_assets", {}) + if not isinstance(asset_name, str) or asset_name not in shared_assets: + raise StarVLAError( + f"variant {variant_entry.get('_catalog_key', variant_entry.get('framework'))!r} " + f"references unknown Qwen asset {asset_name!r}" + ) + return asset_name, dict(shared_assets[asset_name]) + + +def staged_qwen_asset_hashes(qwen_entry: Mapping[str, Any]) -> dict[str, str]: + """Return immutable assets that survive checkpoint surgery unchanged.""" + return { + relative: qwen_entry.get("staged_overrides", {}).get(relative, record)[ + "sha256" + ] + for relative, record in qwen_entry["file_hashes"].items() + if relative not in GENERATED_QWEN_ASSET_PATHS + } + + +def sha256_file(path: Path, chunk_size: int = 8 * 1024 * 1024) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(chunk_size): + digest.update(chunk) + return digest.hexdigest() + + +def verify_checkpoint_file(path: Path, variant_entry: Mapping[str, Any]) -> None: + checkpoint = variant_entry.get("checkpoint") + if checkpoint is None: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no policy checkpoint") + if not path.is_file(): + raise StarVLAError(f"checkpoint does not exist: {path}") + expected_size = int(checkpoint["size"]) + actual_size = path.stat().st_size + if actual_size != expected_size: + raise StarVLAError(f"checkpoint size mismatch for {path}: expected {expected_size}, got {actual_size}") + expected_hash = str(checkpoint["sha256"]) + actual_hash = sha256_file(path) + if actual_hash != expected_hash: + raise StarVLAError(f"checkpoint SHA256 mismatch for {path}: expected {expected_hash}, got {actual_hash}") + + +def verify_catalog_files(root: Path, entry: Mapping[str, Any]) -> dict[str, str]: + verified = {} + for relative in entry.get("files", []): + path = root / relative + expected = entry["file_hashes"][relative] + if not path.is_file(): + raise StarVLAError(f"missing pinned source asset: {path}") + actual_size = path.stat().st_size + actual_hash = sha256_file(path) + if actual_size != expected["size"] or actual_hash != expected["sha256"]: + raise StarVLAError( + f"pinned source asset size/SHA256 mismatch for {path}: " + f"expected {expected['size']}/{expected['sha256']}, got {actual_size}/{actual_hash}" + ) + verified[relative] = actual_hash + return verified + + +def validate_qwen_vlm_destination_names( + base_assets: Path, + qwen_entry: Mapping[str, Any], + records: Sequence[TensorRecord], + *, + backbone: str, +) -> None: + """Bind Qwen2.5 staged tensor names to the pinned canonical HF weight index.""" + if backbone == "qwen3_vl": + return + if backbone != "qwen2_5_vl": + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + index_name = "model.safetensors.index.json" + if index_name not in qwen_entry.get("files", []): + raise StarVLAError("pinned Qwen2.5 asset has no canonical model weight index") + index_record = qwen_entry.get("file_hashes", {}).get(index_name) + if not isinstance(index_record, Mapping): + raise StarVLAError("pinned Qwen2.5 asset has no model weight index hash") + index_path = base_assets / index_name + if not index_path.is_file(): + raise StarVLAError(f"missing pinned Qwen2.5 model weight index: {index_path}") + if ( + index_path.stat().st_size != index_record.get("size") + or sha256_file(index_path) != index_record.get("sha256") + ): + raise StarVLAError( + f"pinned Qwen2.5 model weight index size/SHA256 mismatch: {index_path}" + ) + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError( + f"failed to load pinned Qwen2.5 model weight index {index_path}: {exc}" + ) from exc + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if ( + not isinstance(weight_map, dict) + or not weight_map + or any( + not isinstance(name, str) + or not name + or not isinstance(shard, str) + or not shard + for name, shard in weight_map.items() + ) + ): + raise StarVLAError(f"invalid pinned Qwen2.5 model weight index: {index_path}") + + if any(record.component != "vlm" for record in records): + raise StarVLAError("Qwen VLM destination validation received a non-VLM tensor") + canonical_names = set(weight_map) + staged_backbone_names = { + record.destination_name for record in records if record.role != "lm_head" + } + staged_lm_head_names = { + record.destination_name for record in records if record.role == "lm_head" + } + if staged_lm_head_names != {"lm_head.weight"}: + raise StarVLAError( + "Qwen2.5 staged LM head tensor set mismatch: " + f"expected ['lm_head.weight'], got {sorted(staged_lm_head_names)}" + ) + if staged_backbone_names != canonical_names: + raise StarVLAError( + "Qwen2.5 staged backbone tensor names do not match the pinned canonical " + f"{len(canonical_names)}-tensor HF index; " + f"missing={sorted(canonical_names - staged_backbone_names)[:8]}, " + f"unexpected={sorted(staged_backbone_names - canonical_names)[:8]}" + ) + if len(records) != len(canonical_names) + 1: + raise StarVLAError( + "Qwen2.5 staged VLM tensor count does not equal the canonical HF index " + "plus the checkpoint LM head" + ) + + +def bundle_uuid(variant_entry: Mapping[str, Any], catalog: Mapping[str, Any]) -> str: + """Derive the bundle identity from every source that can change runtime semantics.""" + qwen_asset_name, qwen_entry = get_qwen_asset(catalog, variant_entry) + qwen_hashes = staged_qwen_asset_hashes(qwen_entry) + policy_hashes = { + relative: record["sha256"] for relative, record in variant_entry["file_hashes"].items() + } + provenance = { + "schema_version": 1, + "framework": variant_entry["framework"], + "policy": { + "repo_id": variant_entry["repo_id"], + "revision": variant_entry["revision"], + "checkpoint_sha256": variant_entry["checkpoint"]["sha256"], + "asset_sha256": policy_hashes, + }, + "qwen": { + "repo_id": qwen_entry["repo_id"], + "revision": qwen_entry["revision"], + "staged_asset_sha256": qwen_hashes, + }, + "source_revisions": { + "starvla": catalog["source_revisions"]["starvla"], + "llama_cpp": catalog["source_revisions"]["llama_cpp"], + }, + } + catalog_variant = variant_entry.get("_catalog_key", variant_entry["framework"]) + backbone = variant_entry["backbone"] + if ( + catalog_variant != variant_entry["framework"] + or backbone != "qwen3_vl" + or qwen_asset_name != LEGACY_QWEN3_ASSET + ): + provenance["catalog_variant"] = catalog_variant + provenance["backbone"] = backbone + provenance["qwen"]["asset"] = qwen_asset_name + canonical = json.dumps(provenance, sort_keys=True, separators=(",", ":")) + return str(uuid.uuid5(uuid.NAMESPACE_URL, f"robotcpp:starvla-bundle:{canonical}")) + + +def _set_effective_value( + effective: dict[str, Any], + path: str, + value: Any, +) -> None: + owner: dict[str, Any] = effective + parts = path.split(".") + for part in parts[:-1]: + child = owner.get(part) + if child is None: + child = {} + owner[part] = child + if not isinstance(child, dict): + raise StarVLAError(f"effective config path is not an object: {path}") + owner = child + owner[parts[-1]] = value + + +def resolve_effective_config( + source_dir: Path, + variant_name: str, + variant_entry: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Load the runtime-canonical YAML and apply checkpoint-derived compatibility fixes.""" + framework_name = ( + str(variant_entry["framework"]) if variant_entry is not None else variant_name + ) + backbone = ( + str(variant_entry["backbone"]) + if variant_entry is not None + else "qwen3_vl" + ) + if framework_name not in {"oft", "groot", "pi", "pi_v3"}: + raise StarVLAError( + f"unsupported effective-config variant/framework: " + f"{variant_name!r}/{framework_name!r}" + ) + if backbone not in SUPPORTED_BACKBONES: + raise StarVLAError(f"unsupported effective-config backbone: {backbone!r}") + yaml_path = source_dir / "config.yaml" + if not yaml_path.is_file(): + raise StarVLAError(f"missing canonical StarVLA config: {yaml_path}") + try: + import yaml + except ImportError as exc: + raise StarVLAError("PyYAML is required to resolve the effective StarVLA config") from exc + try: + canonical = yaml.safe_load(yaml_path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise StarVLAError(f"failed to load canonical StarVLA config {yaml_path}: {exc}") from exc + if not isinstance(canonical, dict): + raise StarVLAError(f"expected an object in canonical StarVLA config {yaml_path}") + + effective = deepcopy(canonical) + qwen_hidden_dim = 2048 if backbone == "qwen2_5_vl" else 2560 + if framework_name == "oft": + values = ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.action_hidden_dim", qwen_hidden_dim), + ("framework.action_model.action_model_type", "MLP"), + ) + elif framework_name == "groot": + values = ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.diffusion_model_cfg.cross_attention_dim", qwen_hidden_dim), + ("framework.action_model.diffusion_model_cfg.input_embedding_dim", 768), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ("framework.action_model.diffusion_model_cfg.num_attention_heads", 12), + ) + elif framework_name == "pi_v3": + values = ( + ("framework.qwenvl.vl_hidden_dim", 2560), + ("framework.qwenvl.num_vl_layers", 36), + ("framework.action_model.action_model_type", "LayerwiseFM"), + ("framework.action_model.diffusion_model_cfg.action_dit_hidden_dim", 1024), + ("framework.action_model.diffusion_model_cfg.input_embedding_dim", 1024), + ("framework.action_model.diffusion_model_cfg.cross_attention_dim", 1024), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ("framework.action_model.diffusion_model_cfg.num_attention_heads", 16), + ("framework.action_model.diffusion_model_cfg.num_layers", 36), + ("framework.action_model.diffusion_model_cfg.interleave_self_attention", False), + ("framework.action_model.diffusion_model_cfg.use_canonical_forward", True), + ) + elif framework_name == "pi": + values = ( + ("framework.qwenvl.vl_hidden_dim", qwen_hidden_dim), + ("framework.action_model.hidden_size", qwen_hidden_dim), + ( + "framework.action_model.diffusion_model_cfg.input_embedding_dim", + qwen_hidden_dim, + ), + ( + "framework.action_model.diffusion_model_cfg.cross_attention_dim", + qwen_hidden_dim, + ), + ("framework.action_model.diffusion_model_cfg.attention_head_dim", 64), + ( + "framework.action_model.diffusion_model_cfg.num_attention_heads", + qwen_hidden_dim // 64, + ), + ("framework.action_model.diffusion_model_cfg.use_canonical_forward", False), + ) + else: + raise AssertionError(f"unhandled effective-config framework: {framework_name}") + + for path, value in values: + _set_effective_value(effective, path, value) + _set_effective_value(effective, "framework.action_model.action_horizon", 16) + effective["_robotcpp_effective_config"] = { + "schema_version": 1, + "variant": variant_name, + "framework": framework_name, + "backbone": backbone, + } + return effective + + +def load_checkpoint_state(path: Path) -> dict[str, Any]: + try: + import torch + except ImportError as exc: + raise StarVLAError("PyTorch is required to inspect a StarVLA checkpoint") from exc + + try: + if path.suffix == ".safetensors": + from safetensors.torch import load_file + + raw = load_file(path, device="cpu") + else: + raw = torch.load(path, map_location="cpu", mmap=True, weights_only=True) + except Exception as exc: + raise StarVLAError(f"failed to load checkpoint {path}: {exc}") from exc + + if isinstance(raw, Mapping) and raw and all(isinstance(key, str) and torch.is_tensor(value) for key, value in raw.items()): + return dict(raw) + + if isinstance(raw, Mapping): + for wrapper_key in ("state_dict", "model"): + candidate = raw.get(wrapper_key) + if isinstance(candidate, Mapping) and candidate and all( + isinstance(key, str) and torch.is_tensor(value) for key, value in candidate.items() + ): + unknown_wrappers = set(raw) - {wrapper_key} + if unknown_wrappers: + raise StarVLAError( + f"checkpoint wrapper {wrapper_key!r} has unrecognized sibling keys: {sorted(unknown_wrappers)}" + ) + return dict(candidate) + + raise StarVLAError("checkpoint must be a non-empty flat tensor state_dict or a known single-key wrapper") + + +def classify_tensor(name: str, variant_entry: Mapping[str, Any]) -> tuple[str, str, str]: + backbone = str(variant_entry["backbone"]) + destination_prefixes = VLM_DESTINATION_PREFIXES.get(backbone) + if destination_prefixes is None: + raise StarVLAError(f"unsupported StarVLA Qwen backbone: {backbone!r}") + + for source_prefix, role in VLM_SOURCE_RULES: + if name.startswith(source_prefix): + suffix = name[len(source_prefix) :] + if not suffix: + break + return "vlm", role, destination_prefixes[role] + suffix + + for prefix in variant_entry.get("policy_prefixes", []): + if name.startswith(prefix) and len(name) > len(prefix): + return "policy", "policy", name + if name in variant_entry.get("policy_tensors", []): + return "policy", "policy", name + + raise StarVLAError(f"unrecognized tensor for {variant_entry.get('framework')}: {name}") + + +def _storage_key(tensor: Any) -> tuple[int, int] | None: + try: + storage = tensor.untyped_storage() + return int(storage.data_ptr()), int(storage.nbytes()) + except Exception: + return None + + +def build_inventory( + state_dict: Mapping[str, Any], + variant_entry: Mapping[str, Any], + *, + enforce_expected: bool = True, +) -> list[TensorRecord]: + provisional: list[tuple[TensorRecord, tuple[int, int] | None]] = [] + destinations: dict[str, str] = {} + aliases: dict[tuple[int, int], list[str]] = defaultdict(list) + + for source_name in sorted(state_dict): + tensor = state_dict[source_name] + if not tensor.is_contiguous(): + raise StarVLAError(f"non-contiguous source tensor is not supported: {source_name}") + component, role, destination_name = classify_tensor(source_name, variant_entry) + previous = destinations.get(destination_name) + if previous is not None: + raise StarVLAError( + f"duplicate destination tensor {destination_name!r}: source keys {previous!r} and {source_name!r}" + ) + destinations[destination_name] = source_name + + storage_key = _storage_key(tensor) + if storage_key is not None: + aliases[storage_key].append(source_name) + record = TensorRecord( + source_name=source_name, + destination_name=destination_name, + component=component, + role=role, + shape=[int(dim) for dim in tensor.shape], + dtype=str(tensor.dtype).removeprefix("torch."), + numel=int(tensor.numel()), + nbytes=int(tensor.numel() * tensor.element_size()), + storage_offset=int(tensor.storage_offset()), + ) + provisional.append((record, storage_key)) + + alias_names: dict[tuple[int, int], str] = {} + alias_index = 0 + for storage_key, source_names in sorted(aliases.items(), key=lambda item: min(item[1])): + if len(source_names) > 1: + alias_names[storage_key] = f"alias_{alias_index:04d}" + alias_index += 1 + + records = [ + TensorRecord(**{**record.to_json(), "storage_alias": alias_names.get(storage_key)}) + for record, storage_key in provisional + ] + if enforce_expected: + validate_expected_inventory(records, variant_entry) + return records + + +def inventory_summary(records: list[TensorRecord]) -> dict[str, Any]: + counts = Counter(record.component for record in records) + roles = Counter(record.role for record in records) + numel = Counter() + nbytes = Counter() + dtypes = Counter() + for record in records: + numel[record.component] += record.numel + nbytes[record.component] += record.nbytes + dtypes[record.dtype] += 1 + return { + "total_tensors": len(records), + "vlm_tensors": counts["vlm"], + "policy_tensors": counts["policy"], + "visual_tensors": roles["visual"], + "text_tensors": roles["text"], + "lm_head_tensors": roles["lm_head"], + "total_numel": sum(record.numel for record in records), + "vlm_numel": numel["vlm"], + "policy_numel": numel["policy"], + "total_nbytes": sum(record.nbytes for record in records), + "vlm_nbytes": nbytes["vlm"], + "policy_nbytes": nbytes["policy"], + "dtypes": dict(sorted(dtypes.items())), + "storage_alias_groups": len({record.storage_alias for record in records if record.storage_alias}), + } + + +def validate_expected_inventory(records: list[TensorRecord], variant_entry: Mapping[str, Any]) -> None: + expected = variant_entry.get("expected") + if not expected: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no expected checkpoint inventory") + actual = inventory_summary(records) + mismatches = [] + for key, expected_value in expected.items(): + actual_value = actual.get(key) + if actual_value != expected_value: + mismatches.append(f"{key}: expected {expected_value}, got {actual_value}") + if mismatches: + raise StarVLAError("checkpoint inventory mismatch: " + "; ".join(mismatches)) + + by_destination = {record.destination_name: record for record in records} + shape_mismatches = [] + for name, expected_shape in variant_entry.get("required_shapes", {}).items(): + record = by_destination.get(name) + if record is None: + shape_mismatches.append(f"{name}: missing") + elif record.shape != expected_shape: + shape_mismatches.append(f"{name}: expected {expected_shape}, got {record.shape}") + if shape_mismatches: + raise StarVLAError("checkpoint required-shape mismatch: " + "; ".join(shape_mismatches)) + + +def validate_surgery_manifest( + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], + catalog: Mapping[str, Any], +) -> None: + """Validate a surgery manifest against its catalog entry.""" + checkpoint = variant_entry.get("checkpoint") + if checkpoint is None: + raise StarVLAError(f"variant {variant_entry.get('framework')!r} has no checkpoint") + + expected_top_level = { + "schema_version": 1, + "variant": variant_entry.get("_catalog_key", variant_entry["framework"]), + "model_type": variant_entry["model_type"], + } + _, qwen_entry = get_qwen_asset(catalog, variant_entry) + source = manifest.get("source") + if not isinstance(source, Mapping): + raise StarVLAError("surgery manifest has no source object") + expected_source = { + "repo_id": variant_entry["repo_id"], + "revision": variant_entry["revision"], + "checkpoint_size": checkpoint["size"], + "checkpoint_sha256": checkpoint["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_entry["repo_id"], + "qwen_revision": qwen_entry["revision"], + } + mismatches = [] + for key, expected in expected_top_level.items(): + if manifest.get(key) != expected: + mismatches.append(f"{key}: expected {expected!r}, got {manifest.get(key)!r}") + for key, expected in expected_source.items(): + if source.get(key) != expected: + mismatches.append(f"source.{key}: expected {expected!r}, got {source.get(key)!r}") + + inventory = manifest.get("inventory") + if not isinstance(inventory, Mapping): + mismatches.append("inventory: missing or not an object") + else: + for key, expected in variant_entry.get("expected", {}).items(): + if inventory.get(key) != expected: + mismatches.append(f"inventory.{key}: expected {expected!r}, got {inventory.get(key)!r}") + + expected_uuid = bundle_uuid(variant_entry, catalog) + if manifest.get("bundle_uuid") != expected_uuid: + mismatches.append(f"bundle_uuid: expected {expected_uuid!r}, got {manifest.get('bundle_uuid')!r}") + + qwen_expected = staged_qwen_asset_hashes(qwen_entry) + policy_expected = { + relative: record["sha256"] for relative, record in variant_entry["file_hashes"].items() + } + if manifest.get("qwen_assets") != qwen_expected: + mismatches.append("qwen_assets: staged hashes do not match the pinned Qwen assets") + if manifest.get("policy_assets") != policy_expected: + mismatches.append("policy_assets: staged hashes do not match the pinned policy assets") + + tensors = manifest.get("tensors") + expected_tensor_count = int(variant_entry.get("expected", {}).get("total_tensors", -1)) + if not isinstance(tensors, list) or len(tensors) != expected_tensor_count: + mismatches.append( + f"tensors: expected a {expected_tensor_count}-record inventory, " + f"got {len(tensors) if isinstance(tensors, list) else 'missing'}" + ) + else: + required_record_keys = set(TensorRecord.__dataclass_fields__) + source_names = set() + destination_names = set() + for index, record in enumerate(tensors): + if not isinstance(record, Mapping) or set(record) != required_record_keys: + mismatches.append(f"tensors[{index}]: invalid tensor record schema") + break + source_names.add(record["source_name"]) + destination_names.add(record["destination_name"]) + if len(source_names) != expected_tensor_count or len(destination_names) != expected_tensor_count: + mismatches.append("tensors: source and destination names must be unique") + + for field in ("vlm_output", "policy_output"): + output = manifest.get(field) + if not isinstance(output, Mapping): + mismatches.append(f"{field}: missing or not an object") + continue + if not isinstance(output.get("index"), str) or not isinstance(output.get("shards"), list): + mismatches.append(f"{field}: invalid index/shard records") + if not isinstance(output.get("index_size"), int) or not isinstance(output.get("index_sha256"), str): + mismatches.append(f"{field}: missing index size/SHA256") + + effective_config = manifest.get("effective_config") + if not isinstance(effective_config, Mapping): + mismatches.append("effective_config: missing or not an object") + elif ( + effective_config.get("path") != "effective_config.json" + or not isinstance(effective_config.get("size"), int) + or not isinstance(effective_config.get("sha256"), str) + ): + mismatches.append("effective_config: invalid path/size/SHA256 record") + if mismatches: + raise StarVLAError("surgery manifest does not match the catalog: " + "; ".join(mismatches)) + + +def verify_staged_assets(root: Path, assets: Mapping[str, Any], *, component: str) -> None: + if not isinstance(assets, Mapping) or not assets: + raise StarVLAError(f"surgery manifest has no {component} asset hashes") + for relative, expected_hash in sorted(assets.items()): + relative_path = _safe_relative_path(relative, field=f"{component} assets") + path = root / relative_path + if not path.is_file(): + raise StarVLAError(f"missing staged {component} asset: {path}") + actual_hash = sha256_file(path) + if actual_hash != expected_hash: + raise StarVLAError( + f"staged {component} asset SHA256 mismatch for {path}: expected {expected_hash}, got {actual_hash}" + ) + + +def verify_staged_shards(root: Path, output: Mapping[str, Any], *, component: str) -> dict[str, Any]: + if not isinstance(output, Mapping): + raise StarVLAError(f"surgery manifest has no {component} output object") + index_relative = _safe_relative_path(output.get("index"), field=f"{component} index") + index_path = root / index_relative + if not index_path.is_file(): + raise StarVLAError(f"missing staged {component} index: {index_path}") + expected_index_size = output.get("index_size") + expected_index_hash = output.get("index_sha256") + if index_path.stat().st_size != expected_index_size or sha256_file(index_path) != expected_index_hash: + raise StarVLAError(f"staged {component} index size/SHA256 mismatch: {index_path}") + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load staged {component} index {index_path}: {exc}") from exc + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise StarVLAError(f"invalid or empty staged {component} weight_map: {index_path}") + + shard_records = output.get("shards") + if not isinstance(shard_records, list) or not shard_records: + raise StarVLAError(f"surgery manifest has no {component} shard records") + manifest_shards = set() + tensor_count = 0 + for record in shard_records: + if not isinstance(record, Mapping): + raise StarVLAError(f"invalid {component} shard record: {record!r}") + relative_path = _safe_relative_path(record.get("path"), field=f"{component} shard") + relative = relative_path.as_posix() + if relative in manifest_shards: + raise StarVLAError(f"duplicate {component} shard record: {relative}") + manifest_shards.add(relative) + path = root / relative_path + if not path.is_file(): + raise StarVLAError(f"missing staged {component} shard: {path}") + actual_size = path.stat().st_size + actual_hash = sha256_file(path) + if actual_size != record.get("size") or actual_hash != record.get("sha256"): + raise StarVLAError(f"staged {component} shard size/SHA256 mismatch: {path}") + tensor_count += int(record.get("tensor_count", -1)) + + indexed_shards = {str(value) for value in weight_map.values()} + if indexed_shards != manifest_shards: + raise StarVLAError( + f"staged {component} index/manifest shard mismatch: index={sorted(indexed_shards)}, " + f"manifest={sorted(manifest_shards)}" + ) + if tensor_count != len(weight_map): + raise StarVLAError( + f"staged {component} tensor count mismatch: manifest={tensor_count}, index={len(weight_map)}" + ) + return index + + +def _verify_staged_component( + root: Path, + output: Mapping[str, Any], + state_dict: Mapping[str, Any], + source_records: list[TensorRecord], + *, + component: str, +) -> None: + index = verify_staged_shards(root, output, component=component) + weight_map = index["weight_map"] + records = [record for record in source_records if record.component == component] + expected_names = {record.destination_name for record in records} + if set(weight_map) != expected_names: + raise StarVLAError( + f"staged {component} tensor set does not match the checkpoint: " + f"expected {len(expected_names)}, got {len(weight_map)}" + ) + + try: + import torch + from safetensors import safe_open + except ImportError as exc: + raise StarVLAError("PyTorch and safetensors are required for staged tensor verification") from exc + + by_shard: dict[str, list[TensorRecord]] = defaultdict(list) + for record in records: + by_shard[str(weight_map[record.destination_name])].append(record) + for shard, shard_records in sorted(by_shard.items()): + shard_path = root / _safe_relative_path(shard, field=f"{component} shard index") + with safe_open(shard_path, framework="pt", device="cpu") as handle: + names = {record.destination_name for record in shard_records} + if set(handle.keys()) != names: + raise StarVLAError(f"staged {component} shard/index key mismatch: {shard_path}") + for record in shard_records: + staged = handle.get_tensor(record.destination_name) + original = state_dict[record.source_name].detach().cpu() + if staged.dtype != original.dtype or list(staged.shape) != list(original.shape): + raise StarVLAError( + f"staged tensor dtype/shape mismatch for {record.destination_name}: " + f"expected {original.dtype}/{list(original.shape)}, " + f"got {staged.dtype}/{list(staged.shape)}" + ) + if not torch.equal(staged, original): + raise StarVLAError( + f"staged tensor content does not match the checkpoint: {record.destination_name}" + ) + del staged + + +def verify_staged_components_against_checkpoint( + components: Mapping[str, tuple[Path, Mapping[str, Any]]], + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], +) -> None: + """Bind one or more staged components to a single load of the pinned checkpoint.""" + if not components or not set(components).issubset({"vlm", "policy"}): + raise StarVLAError(f"invalid staged tensor components: {sorted(components)}") + source = manifest.get("source") + if not isinstance(source, Mapping) or not isinstance(source.get("checkpoint"), str): + raise StarVLAError("surgery manifest has no source checkpoint path") + checkpoint_path = Path(source["checkpoint"]) + verify_checkpoint_file(checkpoint_path, variant_entry) + + state_dict = load_checkpoint_state(checkpoint_path) + source_records = build_inventory(state_dict, variant_entry, enforce_expected=True) + manifest_records = manifest.get("tensors") + expected_manifest = [record.to_json() for record in source_records] + if manifest_records != expected_manifest: + raise StarVLAError("surgery tensor inventory does not match the verified checkpoint") + for component, (root, output) in components.items(): + _verify_staged_component( + root, + output, + state_dict, + source_records, + component=component, + ) + del state_dict + + +def verify_staged_tensors_against_checkpoint( + root: Path, + output: Mapping[str, Any], + manifest: Mapping[str, Any], + variant_entry: Mapping[str, Any], + *, + component: str, +) -> None: + """Bind one staged component to the pinned checkpoint.""" + verify_staged_components_against_checkpoint( + {component: (root, output)}, + manifest, + variant_entry, + ) + + +def create_output_temporary(path: Path) -> tuple[int, Path]: + """Create a same-directory temporary file with normal umask-derived permissions.""" + path.parent.mkdir(parents=True, exist_ok=True) + for _ in range(100): + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + descriptor = os.open( + temporary, + os.O_CREAT | os.O_EXCL | os.O_WRONLY, + 0o666, + ) + except FileExistsError: + continue + return descriptor, temporary + raise StarVLAError(f"failed to allocate a temporary output beside {path}") + + +def atomic_write_json(path: Path, value: Any, *, overwrite: bool = True) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = json.dumps(value, indent=2, sort_keys=True) + "\n" + descriptor, temporary = create_output_temporary(path) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + if overwrite: + os.replace(temporary, path) + else: + try: + os.link(temporary, path) + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output: {path}") from exc + temporary.unlink() + finally: + temporary.unlink(missing_ok=True) diff --git a/tools/hf2gguf/starvla/starvla_surgery.py b/tools/hf2gguf/starvla/starvla_surgery.py new file mode 100755 index 0000000..d5d23b3 --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_surgery.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""Split a StarVLA checkpoint into HF-compatible Qwen3-VL and policy staging.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import sys +from pathlib import Path +from typing import Any, Mapping + +from starvla_checkpoint import ( + DEFAULT_CATALOG, + StarVLAError, + TensorRecord, + atomic_write_json, + build_inventory, + get_qwen_asset, + get_variant, + inventory_summary, + load_catalog, + load_checkpoint_state, + bundle_uuid, + resolve_effective_config, + sha256_file, + staged_qwen_asset_hashes, + validate_qwen_vlm_destination_names, + verify_catalog_files, + verify_checkpoint_file, +) + + +def parse_size(value: str) -> int: + match = re.fullmatch(r"([1-9][0-9]*)([KMG]?)", value.strip().upper()) + if not match: + raise argparse.ArgumentTypeError("size must be an integer optionally followed by K, M, or G") + amount = int(match.group(1)) + multiplier = {"": 1, "K": 1024, "M": 1024**2, "G": 1024**3}[match.group(2)] + return amount * multiplier + + +def plan_shards(records: list[TensorRecord], max_shard_size: int) -> list[list[TensorRecord]]: + if max_shard_size <= 0: + raise StarVLAError("max shard size must be positive") + shards: list[list[TensorRecord]] = [] + current: list[TensorRecord] = [] + current_size = 0 + for record in sorted(records, key=lambda item: item.destination_name): + if current and current_size + record.nbytes > max_shard_size: + shards.append(current) + current = [] + current_size = 0 + current.append(record) + current_size += record.nbytes + if current: + shards.append(current) + return shards + + +def _prepare_tensor(tensor: Any, *, clone: bool) -> Any: + prepared = tensor.detach().cpu() + if clone: + prepared = prepared.clone() + elif not prepared.is_contiguous(): + prepared = prepared.contiguous() + return prepared + + +def write_safetensor_shards( + output_dir: Path, + prefix: str, + index_name: str, + records: list[TensorRecord], + state_dict: Mapping[str, Any], + max_shard_size: int, +) -> dict[str, Any]: + try: + from safetensors import safe_open + from safetensors.torch import save_file + except ImportError as exc: + raise StarVLAError("safetensors is required for StarVLA surgery") from exc + + output_dir.mkdir(parents=True, exist_ok=True) + shards = plan_shards(records, max_shard_size) + weight_map: dict[str, str] = {} + shard_manifest = [] + alias_seen: set[str] = set() + + for index, shard_records in enumerate(shards, start=1): + filename = f"{prefix}-{index:05d}-of-{len(shards):05d}.safetensors" + path = output_dir / filename + temporary = output_dir / f".{filename}.tmp" + tensors = {} + for record in shard_records: + clone = bool(record.storage_alias and record.storage_alias in alias_seen) + tensors[record.destination_name] = _prepare_tensor(state_dict[record.source_name], clone=clone) + if record.storage_alias: + alias_seen.add(record.storage_alias) + weight_map[record.destination_name] = filename + + save_file(tensors, temporary, metadata={"format": "pt"}) + os.replace(temporary, path) + with safe_open(path, framework="pt", device="cpu") as handle: + actual_names = set(handle.keys()) + expected_names = set(tensors) + if actual_names != expected_names: + raise StarVLAError(f"safetensors key mismatch after writing {path}") + for record in shard_records: + actual_shape = [int(dim) for dim in handle.get_slice(record.destination_name).get_shape()] + if actual_shape != record.shape: + raise StarVLAError( + f"safetensors shape mismatch for {record.destination_name}: " + f"expected {record.shape}, got {actual_shape}" + ) + shard_manifest.append( + { + "path": filename, + "size": path.stat().st_size, + "sha256": sha256_file(path), + "tensor_count": len(shard_records), + } + ) + del tensors + + expected_destinations = {record.destination_name for record in records} + if set(weight_map) != expected_destinations: + raise StarVLAError("safetensors weight map does not cover every destination tensor exactly once") + index = { + "metadata": {"total_size": sum(record.nbytes for record in records)}, + "weight_map": dict(sorted(weight_map.items())), + } + index_path = output_dir / index_name + atomic_write_json(index_path, index) + return { + "index": index_name, + "index_size": index_path.stat().st_size, + "index_sha256": sha256_file(index_path), + "shards": shard_manifest, + } + + +def copy_qwen_assets(base_assets: Path, hf_dir: Path, asset_entry: Mapping[str, Any]) -> dict[str, str]: + verify_catalog_files(base_assets, asset_entry) + expected_staged_assets = staged_qwen_asset_hashes(asset_entry) + copied = {} + for relative in expected_staged_assets: + source = base_assets / relative + if not source.is_file(): + raise StarVLAError(f"missing pinned Qwen asset: {source}") + destination = hf_dir / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + copied[relative] = sha256_file(destination) + + config_path = hf_dir / "config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["tie_word_embeddings"] = False + text_config = config.get("text_config") + if isinstance(text_config, dict): + text_config["tie_word_embeddings"] = False + elif config.get("model_type") != "qwen2_5_vl": + raise StarVLAError("Qwen-VL config.json has no text_config object") + atomic_write_json(config_path, config) + copied["config.json"] = sha256_file(config_path) + if copied != expected_staged_assets: + raise StarVLAError( + "staged Qwen asset hashes do not match the catalog overrides" + ) + return copied + + +def copy_policy_assets(source_dir: Path, policy_dir: Path, variant_entry: Mapping[str, Any]) -> dict[str, str]: + verify_catalog_files(source_dir, variant_entry) + copied = {} + for relative in variant_entry.get("files", []): + source = source_dir / relative + if not source.is_file(): + raise StarVLAError(f"missing pinned StarVLA policy asset: {source}") + destination = policy_dir / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + copied[relative] = sha256_file(destination) + return copied + + +def _run_surgery_in_owned_directory( + checkpoint: Path, + source_dir: Path, + base_assets: Path, + output_dir: Path, + variant_name: str, + catalog_path: Path, + max_shard_size: int, +) -> dict[str, Any]: + catalog = load_catalog(catalog_path) + variant = get_variant(catalog, variant_name) + if variant.get("checkpoint") is None: + raise StarVLAError(f"variant {variant_name!r} has no policy checkpoint to split") + verify_checkpoint_file(checkpoint, variant) + if output_dir.exists() and any(output_dir.iterdir()): + raise StarVLAError(f"output directory is not empty: {output_dir}") + + state_dict = load_checkpoint_state(checkpoint) + records = build_inventory(state_dict, variant, enforce_expected=True) + vlm_records = [record for record in records if record.component == "vlm"] + policy_records = [record for record in records if record.component == "policy"] + if len(vlm_records) + len(policy_records) != len(records): + raise StarVLAError("source tensor set is not the disjoint union of VLM and policy tensors") + + hf_dir = output_dir / "hf" + policy_dir = output_dir / "policy" + hf_dir.mkdir(parents=True, exist_ok=True) + policy_dir.mkdir(parents=True, exist_ok=True) + + qwen_asset_name, qwen_asset_entry = get_qwen_asset(catalog, variant) + validate_qwen_vlm_destination_names( + base_assets, + qwen_asset_entry, + vlm_records, + backbone=str(variant["backbone"]), + ) + qwen_assets = copy_qwen_assets(base_assets, hf_dir, qwen_asset_entry) + policy_assets = copy_policy_assets(source_dir, policy_dir, variant) + effective_config_path = policy_dir / "effective_config.json" + atomic_write_json( + effective_config_path, + resolve_effective_config(source_dir, variant_name, variant), + ) + effective_config = { + "path": effective_config_path.name, + "size": effective_config_path.stat().st_size, + "sha256": sha256_file(effective_config_path), + } + vlm_output = write_safetensor_shards( + hf_dir, + "model", + "model.safetensors.index.json", + vlm_records, + state_dict, + max_shard_size, + ) + policy_output = write_safetensor_shards( + policy_dir, + "policy", + "policy.safetensors.index.json", + policy_records, + state_dict, + max_shard_size, + ) + + source_uuid = bundle_uuid(variant, catalog) + manifest = { + "schema_version": 1, + "variant": variant_name, + "framework": variant["framework"], + "backbone": variant["backbone"], + "model_type": variant["model_type"], + "bundle_uuid": source_uuid, + "source": { + "repo_id": variant["repo_id"], + "revision": variant["revision"], + "checkpoint": str(checkpoint), + "checkpoint_size": checkpoint.stat().st_size, + "checkpoint_sha256": variant["checkpoint"]["sha256"], + "starvla_revision": catalog["source_revisions"]["starvla"], + "llama_cpp_revision": catalog["source_revisions"]["llama_cpp"], + "qwen_repo_id": qwen_asset_entry["repo_id"], + "qwen_revision": qwen_asset_entry["revision"], + "qwen_asset": qwen_asset_name, + }, + "inventory": inventory_summary(records), + "qwen_assets": qwen_assets, + "policy_assets": policy_assets, + "effective_config": effective_config, + "vlm_output": vlm_output, + "policy_output": policy_output, + "tensors": [record.to_json() for record in records], + } + atomic_write_json(output_dir / "surgery_manifest.json", manifest, overwrite=False) + return manifest + + +def run_surgery( + checkpoint: Path, + source_dir: Path, + base_assets: Path, + output_dir: Path, + variant_name: str, + catalog_path: Path, + max_shard_size: int, +) -> dict[str, Any]: + """Own the staging directory so a failed split cannot poison a retry.""" + output_dir.parent.mkdir(parents=True, exist_ok=True) + try: + output_dir.mkdir() + except FileExistsError as exc: + raise StarVLAError(f"refusing to overwrite existing output directory: {output_dir}") from exc + + try: + return _run_surgery_in_owned_directory( + checkpoint=checkpoint, + source_dir=source_dir, + base_assets=base_assets, + output_dir=output_dir, + variant_name=variant_name, + catalog_path=catalog_path, + max_shard_size=max_shard_size, + ) + except BaseException: + shutil.rmtree(output_dir, ignore_errors=True) + raise + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("checkpoint", type=Path) + parser.add_argument( + "--variant", + required=True, + choices=("oft", "groot", "pi_v3", "qwen25_oft", "qwen25_groot", "qwen25_pi"), + ) + parser.add_argument("--source-dir", type=Path, required=True) + parser.add_argument("--base-assets", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument("--max-shard-size", type=parse_size, default=parse_size("2G")) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + manifest = run_surgery( + checkpoint=args.checkpoint, + source_dir=args.source_dir, + base_assets=args.base_assets, + output_dir=args.output_dir, + variant_name=args.variant, + catalog_path=args.catalog, + max_shard_size=args.max_shard_size, + ) + print(json.dumps(manifest["inventory"], indent=2, sort_keys=True)) + print(f"surgery manifest: {args.output_dir / 'surgery_manifest.json'}") + return 0 + except (StarVLAError, OSError, json.JSONDecodeError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/hf2gguf/starvla/starvla_variant_config.sh b/tools/hf2gguf/starvla/starvla_variant_config.sh new file mode 100644 index 0000000..a4b8076 --- /dev/null +++ b/tools/hf2gguf/starvla/starvla_variant_config.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash + +_STARVLA_CONFIG_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +load_starvla_variant() { + if [[ $# -ne 1 ]]; then + echo "usage: load_starvla_variant VARIANT" >&2 + return 2 + fi + + local config_python="${STARVLA_CONFIG_PYTHON:-python3}" + local assignments + assignments="$("${config_python}" - "${_STARVLA_CONFIG_DIR}" "$1" <<'PY' +import os +import shlex +import sys +from pathlib import Path + +sys.path.insert(0, sys.argv[1]) +from starvla_checkpoint import ( # noqa: E402 + StarVLAError, + artifact_stem, + get_qwen_asset, + get_variant, + load_catalog, +) + +try: + catalog_path = Path(os.environ["STARVLA_CATALOG"]) if "STARVLA_CATALOG" in os.environ else None + catalog = load_catalog(catalog_path) if catalog_path is not None else load_catalog() + variant = get_variant(catalog, sys.argv[2]) + _, qwen = get_qwen_asset(catalog, variant) +except StarVLAError as exc: + print(f"error: {exc}", file=sys.stderr) + raise SystemExit(2) from exc + +checkpoint = variant["checkpoint"] +values = { + "STARVLA_REVISION": catalog["source_revisions"]["starvla"], + "MODEL_TYPE": variant["model_type"], + "FRAMEWORK": variant["framework"], + "CHECKPOINT_REVISION": variant["revision"], + "CHECKPOINT_SHA256": checkpoint["sha256"], + "CHECKPOINT_DIRECTORY": variant["directory"], + "CHECKPOINT_RELATIVE_PATH": checkpoint["path"], + "QWEN_REVISION": qwen["revision"], + "QWEN_DIRECTORY": qwen["directory"], + "ARTIFACT_STEM": artifact_stem(sys.argv[2]), +} +for name, value in values.items(): + print(f"{name}={shlex.quote(str(value))}") +PY + )" || return 2 + eval "${assignments}" +} diff --git a/tools/hf2gguf/starvla/validate_starvla_bundle.py b/tools/hf2gguf/starvla/validate_starvla_bundle.py new file mode 100755 index 0000000..6cf0a80 --- /dev/null +++ b/tools/hf2gguf/starvla/validate_starvla_bundle.py @@ -0,0 +1,1009 @@ +#!/usr/bin/env python3 +"""Validate a converted StarVLA GGUF bundle and write its content manifest.""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from collections import Counter +from pathlib import Path +from typing import Any + +import numpy as np + +from convert_starvla_policy_to_gguf import ( + GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE, + GROOT_TENSOR_MAP, + OFT_ACTION_TOKEN_ID, + OFT_TENSOR_MAP, + PI_SUPPORTED_DIMENSIONS, + PI_TENSOR_MAP, + PI_V3_SUPPORTED_DIMENSIONS, + PI_V3_TENSOR_MAP, + build_groot_metadata, + build_oft_metadata, + build_pi_metadata, + build_pi_v3_metadata, + load_policy_tensors, + resolve_action_token_id, +) +from starvla_checkpoint import ( + DEFAULT_CATALOG, + DEFAULT_MMPROJ_DTYPE, + DEFAULT_POLICY_DTYPE, + DEFAULT_TEXT_DTYPE, + StarVLAError, + atomic_write_json, + get_variant, + load_catalog, + portable_source_record, + sha256_file, + validate_surgery_manifest, + verify_staged_assets, + verify_staged_components_against_checkpoint, +) + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +LLAMA_GGUF_PY = REPOSITORY_ROOT / "third_party" / "llama.cpp" / "gguf-py" +sys.path.insert(0, str(LLAMA_GGUF_PY)) + +try: + import gguf +except ImportError as exc: + raise SystemExit(f"error: failed to import pinned llama.cpp gguf-py: {exc}") from exc + + +def _load_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise StarVLAError(f"failed to load JSON {path}: {exc}") from exc + if not isinstance(value, dict): + raise StarVLAError(f"expected a JSON object in {path}") + return value + + +def field_value(reader: Any, key: str) -> Any: + field = reader.get_field(key) + if field is None: + raise StarVLAError(f"GGUF is missing required metadata: {key}") + return field.contents() + + +def expect_field(reader: Any, key: str, expected: Any) -> None: + actual = field_value(reader, key) + if actual != expected: + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected {expected!r}, got {actual!r}") + + +def expect_sequence_field(reader: Any, key: str, expected: list[Any]) -> None: + actual = field_value(reader, key) + if not isinstance(actual, list): + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected an array, got {type(actual).__name__}") + if len(actual) != len(expected): + raise StarVLAError( + f"GGUF metadata length mismatch for {key}: expected {len(expected)}, got {len(actual)}" + ) + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + if actual_item != expected_item: + raise StarVLAError( + f"GGUF metadata mismatch for {key}[{index}]: " + f"expected {expected_item!r}, got {actual_item!r}" + ) + + +def tensor_map(reader: Any) -> dict[str, Any]: + tensors = {tensor.name: tensor for tensor in reader.tensors} + if len(tensors) != len(reader.tensors): + raise StarVLAError("GGUF contains duplicate tensor names") + return tensors + + +def expect_ggml_tensor_shape(tensors: dict[str, Any], name: str, shape: list[int]) -> None: + """Check GGUF/ggml dimensions (`ne[]` order), not NumPy/PyTorch dimensions.""" + tensor = tensors.get(name) + if tensor is None: + raise StarVLAError(f"GGUF is missing required tensor: {name}") + actual = [int(dim) for dim in tensor.shape] + if actual != shape: + raise StarVLAError(f"GGUF tensor shape mismatch for {name}: expected {shape}, got {actual}") + + +def expect_complete_tensor_map(tensors: dict[str, Any], expected: dict[str, list[int]], component: str) -> None: + actual_names = set(tensors) + expected_names = set(expected) + if actual_names != expected_names: + raise StarVLAError( + f"{component} GGUF tensor set mismatch; " + f"missing={sorted(expected_names - actual_names)}, unexpected={sorted(actual_names - expected_names)}" + ) + for name, shape in expected.items(): + expect_ggml_tensor_shape(tensors, name, shape) + + +def expected_text_tensor_map( + backbone: str = "qwen3_vl", vocab_size: int = 151936 +) -> dict[str, list[int]]: + if backbone == "qwen2_5_vl": + expected = { + "token_embd.weight": [2048, vocab_size], + "output_norm.weight": [2048], + "output.weight": [2048, vocab_size], + } + per_block = { + "attn_norm.weight": [2048], + "ffn_norm.weight": [2048], + "attn_q.weight": [2048, 2048], + "attn_q.bias": [2048], + "attn_k.weight": [2048, 256], + "attn_k.bias": [256], + "attn_v.weight": [2048, 256], + "attn_v.bias": [256], + "attn_output.weight": [2048, 2048], + "ffn_gate.weight": [2048, 11008], + "ffn_up.weight": [2048, 11008], + "ffn_down.weight": [11008, 2048], + } + for block in range(36): + for suffix, shape in per_block.items(): + expected[f"blk.{block}.{suffix}"] = shape + return expected + if backbone != "qwen3_vl": + raise StarVLAError(f"unsupported Qwen text tensor backbone: {backbone!r}") + expected = { + "token_embd.weight": [2560, 151936], + "output_norm.weight": [2560], + "output.weight": [2560, 151936], + } + per_block = { + "attn_norm.weight": [2560], + "ffn_norm.weight": [2560], + "attn_q.weight": [2560, 4096], + "attn_k.weight": [2560, 1024], + "attn_v.weight": [2560, 1024], + "attn_output.weight": [4096, 2560], + "attn_q_norm.weight": [128], + "attn_k_norm.weight": [128], + "ffn_gate.weight": [2560, 9728], + "ffn_up.weight": [2560, 9728], + "ffn_down.weight": [9728, 2560], + } + for block in range(36): + for suffix, shape in per_block.items(): + expected[f"blk.{block}.{suffix}"] = shape + return expected + + +def expected_mmproj_tensor_map( + backbone: str = "qwen3_vl", +) -> dict[str, list[int]]: + if backbone == "qwen2_5_vl": + expected = { + "v.patch_embd.weight": [14, 14, 3, 1280], + "v.patch_embd.weight.1": [14, 14, 3, 1280], + "v.post_ln.weight": [1280], + "mm.0.weight": [5120, 5120], + "mm.0.bias": [5120], + "mm.2.weight": [5120, 2048], + "mm.2.bias": [2048], + } + per_block = { + "ln1.weight": [1280], + "ln2.weight": [1280], + "attn_q.weight": [1280, 1280], + "attn_q.bias": [1280], + "attn_k.weight": [1280, 1280], + "attn_k.bias": [1280], + "attn_v.weight": [1280, 1280], + "attn_v.bias": [1280], + "attn_out.weight": [1280, 1280], + "attn_out.bias": [1280], + "ffn_gate.weight": [1280, 3420], + "ffn_gate.bias": [3420], + "ffn_up.weight": [1280, 3420], + "ffn_up.bias": [3420], + "ffn_down.weight": [3420, 1280], + "ffn_down.bias": [1280], + } + for block in range(32): + for suffix, shape in per_block.items(): + expected[f"v.blk.{block}.{suffix}"] = shape + return expected + if backbone != "qwen3_vl": + raise StarVLAError(f"unsupported Qwen mmproj tensor backbone: {backbone!r}") + expected = { + "v.position_embd.weight": [1024, 2304], + "v.patch_embd.weight": [16, 16, 3, 1024], + "v.patch_embd.weight.1": [16, 16, 3, 1024], + "v.patch_embd.bias": [1024], + "v.post_ln.weight": [1024], + "v.post_ln.bias": [1024], + "mm.0.weight": [4096, 4096], + "mm.0.bias": [4096], + "mm.2.weight": [4096, 2560], + "mm.2.bias": [2560], + } + per_block = { + "attn_out.weight": [1024, 1024], + "attn_out.bias": [1024], + "attn_qkv.weight": [1024, 3072], + "attn_qkv.bias": [3072], + "ffn_up.weight": [1024, 4096], + "ffn_up.bias": [4096], + "ffn_down.weight": [4096, 1024], + "ffn_down.bias": [1024], + "ln1.weight": [1024], + "ln1.bias": [1024], + "ln2.weight": [1024], + "ln2.bias": [1024], + } + for block in range(24): + for suffix, shape in per_block.items(): + expected[f"v.blk.{block}.{suffix}"] = shape + for layer in (5, 11, 17): + expected.update( + { + f"v.deepstack.{layer}.norm.weight": [4096], + f"v.deepstack.{layer}.norm.bias": [4096], + f"v.deepstack.{layer}.fc1.weight": [4096, 4096], + f"v.deepstack.{layer}.fc1.bias": [4096], + f"v.deepstack.{layer}.fc2.weight": [4096, 2560], + f"v.deepstack.{layer}.fc2.bias": [2560], + } + ) + return expected + + +def metadata_matches(actual: Any, expected: Any) -> bool: + if isinstance(expected, float): + return isinstance(actual, (int, float)) and math.isclose(actual, expected, rel_tol=1e-6, abs_tol=1e-6) + if isinstance(expected, list): + return isinstance(actual, list) and len(actual) == len(expected) and all( + metadata_matches(actual_item, expected_item) + for actual_item, expected_item in zip(actual, expected) + ) + return actual == expected + + +def expect_metadata_field(reader: Any, key: str, expected: Any) -> None: + actual = field_value(reader, key) + if not metadata_matches(actual, expected): + raise StarVLAError(f"GGUF metadata mismatch for {key}: expected {expected!r}, got {actual!r}") + + +def _require_object(value: Any, field: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise StarVLAError(f"pinned Qwen {field} must be a JSON object") + return value + + +def _require_positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise StarVLAError(f"pinned Qwen {field} must be a positive integer") + return value + + +def _special_token_content(value: Any, field: str) -> str | None: + if value is None: + return None + if isinstance(value, str): + return value + if isinstance(value, dict) and isinstance(value.get("content"), str): + return str(value["content"]) + raise StarVLAError(f"pinned Qwen tokenizer_config.json {field} has an unsupported value") + + +def _normalize_merge(merge: Any, index: int) -> str: + if isinstance(merge, str): + return merge + if ( + isinstance(merge, list) + and len(merge) == 2 + and all(isinstance(part, str) for part in merge) + ): + encoded = [ + "".join(chr(ord(character) + 256) if character == " " else character for character in part) + for part in merge + ] + return " ".join(encoded) + raise StarVLAError(f"pinned Qwen tokenizer merge {index} has an unsupported value") + + +def expected_tokenizer_metadata(hf_dir: Path) -> dict[str, Any]: + """Derive llama.cpp's GPT-2 vocabulary metadata from the pinned HF tokenizer files.""" + tokenizer = _load_json(hf_dir / "tokenizer.json") + tokenizer_config = _load_json(hf_dir / "tokenizer_config.json") + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vocab_size = _require_positive_int(text_config.get("vocab_size"), "text_config.vocab_size") + + model = _require_object(tokenizer.get("model"), "tokenizer.json model") + vocabulary = _require_object(model.get("vocab"), "tokenizer.json model.vocab") + added_tokens = tokenizer.get("added_tokens") + if not isinstance(added_tokens, list): + raise StarVLAError("pinned Qwen tokenizer.json added_tokens must be an array") + + decoder = tokenizer_config.get("added_tokens_decoder") + if not isinstance(decoder, dict): + raise StarVLAError("pinned Qwen tokenizer_config.json added_tokens_decoder must be an object") + decoder_by_id: dict[int, dict[str, Any]] = {} + for raw_id, record in decoder.items(): + if not isinstance(raw_id, str) or not raw_id.isdecimal() or not isinstance(record, dict): + raise StarVLAError("pinned Qwen added_tokens_decoder contains an invalid entry") + token_id = int(raw_id) + if token_id in decoder_by_id: + raise StarVLAError(f"pinned Qwen added_tokens_decoder repeats token id {token_id}") + decoder_by_id[token_id] = record + + tokens = [f"[PAD{token_id}]" for token_id in range(vocab_size)] + token_types = [int(gguf.TokenType.UNUSED)] * vocab_size + assigned_ids: set[int] = set() + token_to_id: dict[str, int] = {} + + def assign(token: Any, token_id: Any, token_type: Any, source: str) -> None: + if not isinstance(token, str): + raise StarVLAError(f"pinned Qwen {source} token must be a string") + if isinstance(token_id, bool) or not isinstance(token_id, int) or not 0 <= token_id < vocab_size: + raise StarVLAError(f"pinned Qwen {source} token id is out of range: {token_id!r}") + if token_id in assigned_ids: + raise StarVLAError(f"pinned Qwen tokenizer repeats token id {token_id}") + if token in token_to_id: + raise StarVLAError(f"pinned Qwen tokenizer repeats token content {token!r}") + assigned_ids.add(token_id) + token_to_id[token] = token_id + tokens[token_id] = token + token_types[token_id] = int(token_type) + + for token, token_id in vocabulary.items(): + assign(token, token_id, gguf.TokenType.NORMAL, "base vocabulary") + + added_by_id: dict[int, dict[str, Any]] = {} + for index, record in enumerate(added_tokens): + if not isinstance(record, dict): + raise StarVLAError(f"pinned Qwen tokenizer added token {index} must be an object") + token_id = record.get("id") + if isinstance(token_id, bool) or not isinstance(token_id, int): + raise StarVLAError(f"pinned Qwen tokenizer added token {index} has an invalid id") + if token_id in added_by_id: + raise StarVLAError(f"pinned Qwen tokenizer repeats added token id {token_id}") + added_by_id[token_id] = record + if set(added_by_id) != set(decoder_by_id): + raise StarVLAError("pinned Qwen tokenizer added_tokens and added_tokens_decoder ids differ") + + for token_id, record in sorted(added_by_id.items()): + decoder_record = decoder_by_id[token_id] + for field in ("content", "normalized", "special"): + if record.get(field) != decoder_record.get(field): + raise StarVLAError( + f"pinned Qwen added token {token_id} disagrees with added_tokens_decoder for {field}" + ) + token = record.get("content") + if not isinstance(token, str): + raise StarVLAError(f"pinned Qwen added token {token_id} has invalid content") + is_control = bool(record.get("special")) or (token.startswith("<|") and token.endswith("|>")) + token_type = gguf.TokenType.CONTROL if is_control else gguf.TokenType.USER_DEFINED + assign(token, token_id, token_type, "added vocabulary") + + raw_merges = model.get("merges") + if not isinstance(raw_merges, list) or not raw_merges: + raise StarVLAError("pinned Qwen tokenizer.json model.merges must be a non-empty array") + merges = [_normalize_merge(merge, index) for index, merge in enumerate(raw_merges)] + + chat_template = tokenizer_config.get("chat_template") + if not isinstance(chat_template, str) or not chat_template: + jinja_path = hf_dir / "chat_template.jinja" + if not jinja_path.is_file(): + raise StarVLAError( + "pinned Qwen tokenizer has no non-empty chat template" + ) + chat_template = jinja_path.read_text(encoding="utf-8") + if not chat_template: + raise StarVLAError("pinned Qwen chat_template.jinja is empty") + chat_template_path = hf_dir / "chat_template.json" + if chat_template_path.is_file(): + template_file = _load_json(chat_template_path).get("chat_template") + if template_file != chat_template: + raise StarVLAError("pinned Qwen chat_template.json disagrees with tokenizer_config.json") + + bos_id = _require_positive_int(text_config.get("bos_token_id"), "text_config.bos_token_id") + eos_id = _require_positive_int(text_config.get("eos_token_id"), "text_config.eos_token_id") + if bos_id >= vocab_size or eos_id >= vocab_size: + raise StarVLAError("pinned Qwen BOS/EOS token id is outside the configured vocabulary") + eos_content = _special_token_content(tokenizer_config.get("eos_token"), "eos_token") + if eos_content is not None and token_to_id.get(eos_content) != eos_id: + raise StarVLAError("pinned Qwen EOS token string and id disagree") + pad_content = _special_token_content(tokenizer_config.get("pad_token"), "pad_token") + if pad_content is None or pad_content not in token_to_id: + raise StarVLAError("pinned Qwen tokenizer has no resolvable padding token") + add_bos = tokenizer_config.get("add_bos_token") + if not isinstance(add_bos, bool): + raise StarVLAError("pinned Qwen tokenizer_config.json add_bos_token must be boolean") + + return { + "tokenizer.ggml.model": "gpt2", + "tokenizer.ggml.pre": "qwen2", + "tokenizer.ggml.tokens": tokens, + "tokenizer.ggml.token_type": token_types, + "tokenizer.ggml.merges": merges, + "tokenizer.ggml.bos_token_id": bos_id, + "tokenizer.ggml.eos_token_id": eos_id, + "tokenizer.ggml.padding_token_id": token_to_id[pad_content], + "tokenizer.ggml.add_bos_token": add_bos, + "tokenizer.chat_template": chat_template, + } + + +def validate_policy_metadata(reader: Any, expected: dict[str, Any]) -> None: + expected_keys = {key for key in expected if key.startswith("starvla.")} + actual_keys = {key for key in reader.fields if key.startswith("starvla.")} + if actual_keys != expected_keys: + raise StarVLAError( + "policy GGUF StarVLA metadata set mismatch; " + f"missing={sorted(expected_keys - actual_keys)}, unexpected={sorted(actual_keys - expected_keys)}" + ) + for key in sorted(expected_keys): + actual = field_value(reader, key) + if not metadata_matches(actual, expected[key]): + raise StarVLAError( + f"policy GGUF metadata mismatch for {key}: expected {expected[key]!r}, got {actual!r}" + ) + expect_field(reader, "general.name", expected["general.name"]) + + +def validate_qwen_vl_image_metadata( + reader: Any, expected_metadata: dict[str, Any], backbone: str +) -> None: + if backbone not in ("qwen3_vl", "qwen2_5_vl"): + raise StarVLAError(f"unsupported Qwen image metadata backbone: {backbone!r}") + expected = { + key: value + for key, value in expected_metadata.items() + if key.startswith("starvla.image.") + } + actual_keys = { + key for key in reader.fields if key.startswith("starvla.image.") + } + if actual_keys != set(expected): + raise StarVLAError( + "Qwen-VL image metadata set mismatch; " + f"missing={sorted(set(expected) - actual_keys)}, " + f"unexpected={sorted(actual_keys - set(expected))}" + ) + for key, value in sorted(expected.items()): + expect_metadata_field(reader, key, value) + + +def validate_dtype_set(reader: Any, requested: str, *, component: str, exact: bool = False) -> dict[str, int]: + requested_type = { + "f32": "F32", + "fp32": "F32", + "f16": "F16", + "bf16": "BF16", + "q8_0": "Q8_0", + }[requested] + counts = Counter(tensor.tensor_type.name for tensor in reader.tensors) + allowed = {requested_type} if exact or requested_type == "F32" else {requested_type, "F32"} + unexpected = set(counts) - allowed + if unexpected or requested_type not in counts: + raise StarVLAError( + f"unexpected {component} GGUF tensor dtypes for {requested}: " + f"counts={dict(sorted(counts.items()))}, allowed={sorted(allowed)}" + ) + return dict(sorted(counts.items())) + + +def _convert_policy_tensor_data(tensor: Any, dtype: str) -> np.ndarray: + array = np.asarray(tensor.detach().float().cpu().numpy(), dtype=np.float32) + if dtype == "fp32": + return array + if dtype == "f16": + return array.astype(np.float16) + if dtype == "bf16": + return gguf.quantize(array, gguf.GGMLQuantizationType.BF16) + raise StarVLAError(f"unsupported policy GGUF dtype: {dtype}") + + +def _first_byte_mismatch(actual: np.ndarray, expected: np.ndarray) -> int | None: + actual_bytes = np.ascontiguousarray(actual).view(np.uint8).reshape(-1) + expected_bytes = np.ascontiguousarray(expected).view(np.uint8).reshape(-1) + if actual_bytes.size != expected_bytes.size: + return min(actual_bytes.size, expected_bytes.size) + chunk_size = 16 * 1024 * 1024 + for offset in range(0, actual_bytes.size, chunk_size): + stop = min(offset + chunk_size, actual_bytes.size) + actual_chunk = actual_bytes[offset:stop] + expected_chunk = expected_bytes[offset:stop] + if not np.array_equal(actual_chunk, expected_chunk): + mismatch = np.flatnonzero(actual_chunk != expected_chunk) + return offset + int(mismatch[0]) + return None + + +def validate_policy_tensor_bytes( + tensors: dict[str, Any], + policy_dir: Path, + dtype: str, + tensor_name_map: dict[str, str] | None = None, + component_label: str = "OFT", +) -> None: + tensor_name_map = OFT_TENSOR_MAP if tensor_name_map is None else tensor_name_map + source_tensors = load_policy_tensors(policy_dir) + missing = sorted(set(tensor_name_map) - set(source_tensors)) + if missing: + raise StarVLAError( + f"staged {component_label} policy is missing runtime tensors: {missing}" + ) + if set(tensors) != set(tensor_name_map.values()): + raise StarVLAError( + f"{component_label} policy GGUF tensor names do not match the canonical " + f"{len(tensor_name_map)}-tensor map" + ) + for source_name, destination_name in tensor_name_map.items(): + source = source_tensors[source_name] + expected_shape = list(reversed(source.shape)) + tensor = tensors[destination_name] + actual_shape = [int(dimension) for dimension in tensor.shape] + if actual_shape != expected_shape: + raise StarVLAError( + f"policy GGUF tensor shape mismatch for {destination_name}: " + f"expected {expected_shape}, got {actual_shape}" + ) + expected = _convert_policy_tensor_data(source, dtype) + actual = np.asarray(tensor.data) + if actual.nbytes != expected.nbytes: + raise StarVLAError( + f"policy GGUF tensor byte size mismatch for {destination_name}: " + f"expected {expected.nbytes}, got {actual.nbytes}" + ) + mismatch = _first_byte_mismatch(actual, expected) + if mismatch is not None: + raise StarVLAError( + f"policy GGUF tensor content mismatch for {destination_name} at byte offset {mismatch}" + ) + del expected + del source_tensors + + +def validate_text( + reader: Any, + bundle_uuid: str, + dtype: str, + hf_dir: Path, + *, + require_oft_action_token: bool = True, + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + if backbone == "qwen3_vl": + architecture = "qwen3vl" + expect_field(reader, "general.architecture", architecture) + expect_field(reader, "qwen3vl.context_length", 262144) + expect_field(reader, "qwen3vl.embedding_length", 2560) + expect_field(reader, "qwen3vl.feed_forward_length", 9728) + expect_field(reader, "qwen3vl.block_count", 36) + expect_field(reader, "qwen3vl.attention.head_count", 32) + expect_field(reader, "qwen3vl.attention.head_count_kv", 8) + expect_field(reader, "qwen3vl.attention.key_length", 128) + expect_field(reader, "qwen3vl.attention.value_length", 128) + expect_metadata_field( + reader, "qwen3vl.attention.layer_norm_rms_epsilon", 1e-6 + ) + expect_field(reader, "qwen3vl.rope.dimension_sections", [24, 20, 20, 0]) + expect_metadata_field(reader, "qwen3vl.rope.freq_base", 5_000_000.0) + expect_field(reader, "qwen3vl.n_deepstack_layers", 3) + vocab_size = 151936 + elif backbone == "qwen2_5_vl": + architecture = "qwen2vl" + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vocab_size = _require_positive_int( + text_config.get("vocab_size"), "text_config.vocab_size" + ) + expect_field(reader, "general.architecture", architecture) + expect_field(reader, "qwen2vl.context_length", 128000) + expect_field(reader, "qwen2vl.embedding_length", 2048) + expect_field(reader, "qwen2vl.feed_forward_length", 11008) + expect_field(reader, "qwen2vl.block_count", 36) + expect_field(reader, "qwen2vl.attention.head_count", 16) + expect_field(reader, "qwen2vl.attention.head_count_kv", 2) + expect_metadata_field( + reader, "qwen2vl.attention.layer_norm_rms_epsilon", 1e-6 + ) + expect_field(reader, "qwen2vl.rope.dimension_sections", [16, 24, 24, 0]) + expect_metadata_field(reader, "qwen2vl.rope.freq_base", 1_000_000.0) + if "qwen2vl.n_deepstack_layers" in reader.fields: + raise StarVLAError("Qwen2.5-VL text GGUF unexpectedly enables DeepStack") + else: + raise StarVLAError(f"unsupported Qwen text backbone: {backbone!r}") + tokenizer_metadata = expected_tokenizer_metadata(hf_dir) + for key, expected in tokenizer_metadata.items(): + if isinstance(expected, list): + expect_sequence_field(reader, key, expected) + else: + expect_field(reader, key, expected) + if require_oft_action_token: + action_token_id = resolve_action_token_id(hf_dir) + if action_token_id != OFT_ACTION_TOKEN_ID: + raise StarVLAError( + f"Qwen action token id mismatch: expected {OFT_ACTION_TOKEN_ID}, got {action_token_id}" + ) + tensors = tensor_map(reader) + expect_complete_tensor_map( + tensors, + expected_text_tensor_map(backbone, vocab_size), + "Qwen text", + ) + return { + "architecture": architecture, + "tensor_count": len(tensors), + "dtypes": validate_dtype_set(reader, dtype, component="text"), + } + + +def validate_mmproj( + reader: Any, + bundle_uuid: str, + dtype: str, + hf_dir: Path, + *, + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + config = _load_json(hf_dir / "config.json") + text_config_value = config.get("text_config") + text_config = ( + _require_object(text_config_value, "config.json text_config") + if text_config_value is not None + else config + ) + vision_config = _require_object(config.get("vision_config"), "config.json vision_config") + preprocessor = _load_json(hf_dir / "preprocessor_config.json") + patch_size = _require_positive_int(vision_config.get("patch_size"), "vision_config.patch_size") + image_mean = preprocessor.get("image_mean") + image_std = preprocessor.get("image_std") + if not isinstance(image_mean, list) or not isinstance(image_std, list): + raise StarVLAError("pinned Qwen preprocessor image_mean/image_std must be arrays") + + expect_field(reader, "general.architecture", "clip") + expect_field(reader, "general.source.uuid", bundle_uuid) + expect_field(reader, "clip.has_vision_encoder", True) + expect_field(reader, "clip.vision.patch_size", patch_size) + expect_field(reader, "clip.vision.embedding_length", vision_config.get("hidden_size")) + expect_field(reader, "clip.vision.feed_forward_length", vision_config.get("intermediate_size")) + expect_field(reader, "clip.vision.projection_dim", text_config.get("hidden_size")) + expect_field(reader, "clip.vision.block_count", vision_config.get("depth")) + expect_field(reader, "clip.vision.attention.head_count", vision_config.get("num_heads")) + expect_metadata_field(reader, "clip.vision.attention.layer_norm_epsilon", text_config.get("rms_norm_eps")) + expect_metadata_field(reader, "clip.vision.image_mean", image_mean) + expect_metadata_field(reader, "clip.vision.image_std", image_std) + if backbone == "qwen3_vl": + num_positions = _require_positive_int( + vision_config.get("num_position_embeddings"), + "vision_config.num_position_embeddings", + ) + positions_per_side = math.isqrt(num_positions) + if positions_per_side * positions_per_side != num_positions: + raise StarVLAError("pinned Qwen vision position count is not square") + deepstack_indices = vision_config.get("deepstack_visual_indexes") + if not isinstance(deepstack_indices, list) or any( + isinstance(index, bool) or not isinstance(index, int) + for index in deepstack_indices + ): + raise StarVLAError( + "pinned Qwen vision_config.deepstack_visual_indexes must be an integer array" + ) + expect_field(reader, "clip.projector_type", "qwen3vl_merger") + expect_field( + reader, "clip.vision.image_size", positions_per_side * patch_size + ) + expect_field(reader, "clip.use_gelu", True) + expect_field( + reader, + "clip.vision.spatial_merge_size", + vision_config.get("spatial_merge_size"), + ) + deepstack = field_value(reader, "clip.vision.is_deepstack_layers") + if ( + len(deepstack) != int(vision_config.get("depth", -1)) + or [ + index for index, enabled in enumerate(deepstack) if enabled + ] + != deepstack_indices + ): + raise StarVLAError( + f"Qwen vision DeepStack layer mismatch: {deepstack}" + ) + elif backbone == "qwen2_5_vl": + expect_field(reader, "clip.projector_type", "qwen2.5vl_merger") + expect_field(reader, "clip.vision.image_size", 560) + expect_field(reader, "clip.use_silu", True) + expect_field(reader, "clip.vision.n_wa_pattern", 8) + if "clip.vision.is_deepstack_layers" in reader.fields: + raise StarVLAError( + "Qwen2.5-VL mmproj unexpectedly contains DeepStack metadata" + ) + else: + raise StarVLAError(f"unsupported Qwen mmproj backbone: {backbone!r}") + tensors = tensor_map(reader) + expect_complete_tensor_map( + tensors, expected_mmproj_tensor_map(backbone), "Qwen mmproj" + ) + return { + "architecture": "clip", + "tensor_count": len(tensors), + "dtypes": validate_dtype_set(reader, dtype, component="mmproj"), + } + + +def validate_policy( + reader: Any, + bundle_uuid: str, + dtype: str, + policy_dir: Path, + text_filename: str, + mmproj_filename: str, + expected_metadata: dict[str, Any], + framework: str = "oft", + backbone: str = "qwen3_vl", +) -> dict[str, Any]: + if framework not in ("oft", "groot", "pi", "pi_v3"): + raise StarVLAError(f"unsupported policy framework for validation: {framework}") + model_type = str(expected_metadata["starvla.model_type"]) + expect_field(reader, "general.architecture", "starvla-policy") + expect_field(reader, "general.source.uuid", bundle_uuid) + expect_field(reader, "starvla.bundle.uuid", bundle_uuid) + expect_field(reader, "starvla.framework", framework) + expect_field(reader, "starvla.model_type", model_type) + expect_field(reader, "starvla.component.text.filename", text_filename) + expect_field(reader, "starvla.component.mmproj.filename", mmproj_filename) + expect_field(reader, "starvla.backbone.arch", backbone) + if framework == "oft": + expect_field(reader, "starvla.prompt.action_token_id", 146663) + elif framework == "groot": + expect_field(reader, "starvla.groot.timestep_ids", [0, 250, 500, 750]) + elif framework == "pi": + expect_field( + reader, + "starvla.conditioning.hidden_tuple_indices", + list(range(21, 37)), + ) + expect_field(reader, "starvla.pi.timestep_ids", [0, 250, 500, 750]) + expect_field( + reader, "starvla.image.framework_inference_pre_resize_width", 224 + ) + expect_field( + reader, "starvla.image.framework_inference_pre_resize_height", 224 + ) + validate_qwen_vl_image_metadata(reader, expected_metadata, backbone) + expect_field(reader, "starvla.action.dimension", 7) + expect_field(reader, "starvla.action.horizon", 16) + expect_field(reader, "starvla.normalization.profile_count", 2) + expect_field( + reader, + "starvla.normalization.profile_keys", + expected_metadata["starvla.normalization.profile_keys"], + ) + validate_policy_metadata(reader, expected_metadata) + tensors = tensor_map(reader) + tensor_name_map = { + "oft": OFT_TENSOR_MAP, + "groot": GROOT_TENSOR_MAP, + "pi": PI_TENSOR_MAP, + "pi_v3": PI_V3_TENSOR_MAP, + }[framework] + dtype_counts = validate_dtype_set(reader, dtype, component="policy", exact=True) + validate_policy_tensor_bytes( + tensors, + policy_dir, + dtype, + tensor_name_map=tensor_name_map, + component_label=framework.upper(), + ) + return { + "architecture": "starvla-policy", + "tensor_count": len(tensors), + "dtypes": dtype_counts, + } + + +def component_record(path: Path, validation: dict[str, Any]) -> dict[str, Any]: + return { + "filename": path.name, + "size": path.stat().st_size, + "sha256": sha256_file(path), + **validation, + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + default="oft", + choices=( + "oft", + "groot", + "pi_v3", + "qwen25_oft", + "qwen25_groot", + "qwen25_pi", + ), + ) + parser.add_argument("--text", type=Path, required=True) + parser.add_argument("--mmproj", type=Path, required=True) + parser.add_argument("--policy", type=Path, required=True) + parser.add_argument("--hf-dir", type=Path, required=True) + parser.add_argument("--policy-dir", type=Path, required=True) + parser.add_argument("--surgery-manifest", type=Path, required=True) + parser.add_argument("--catalog", type=Path, default=DEFAULT_CATALOG) + parser.add_argument( + "--text-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_TEXT_DTYPE, + ) + parser.add_argument( + "--mmproj-dtype", + choices=("f32", "f16", "bf16", "q8_0"), + default=DEFAULT_MMPROJ_DTYPE, + ) + parser.add_argument( + "--policy-dtype", + choices=("fp32", "f16", "bf16"), + default=DEFAULT_POLICY_DTYPE, + ) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + if args.output.exists() or args.output.is_symlink(): + raise StarVLAError(f"refusing to overwrite existing output: {args.output}") + for path in (args.text, args.mmproj, args.policy): + if not path.is_file() or path.stat().st_size == 0: + raise StarVLAError(f"missing or empty bundle component: {path}") + surgery_manifest = _load_json(args.surgery_manifest) + catalog = load_catalog(args.catalog) + variant = get_variant(catalog, args.variant) + framework = str(variant["framework"]) + backbone = str(variant["backbone"]) + validate_surgery_manifest(surgery_manifest, variant, catalog) + verify_staged_assets(args.hf_dir, surgery_manifest.get("qwen_assets", {}), component="Qwen") + verify_staged_assets(args.policy_dir, surgery_manifest.get("policy_assets", {}), component="policy") + verify_staged_components_against_checkpoint( + { + "vlm": (args.hf_dir, surgery_manifest.get("vlm_output", {})), + "policy": (args.policy_dir, surgery_manifest.get("policy_output", {})), + }, + surgery_manifest, + variant, + ) + bundle_uuid = str(surgery_manifest["bundle_uuid"]) + if framework == "oft": + oft_dimensions = ( + {"input_dim": 2048, "hidden_dim": 4096, "action_dim": 7} + if backbone == "qwen2_5_vl" + else {"input_dim": 2560, "hidden_dim": 5120, "action_dim": 7} + ) + expected_policy_metadata = build_oft_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + oft_dimensions, + OFT_ACTION_TOKEN_ID, + args.text.name, + args.mmproj.name, + ) + elif framework == "groot": + groot_dimensions = dict( + GROOT_SUPPORTED_DIMENSIONS_BY_BACKBONE[backbone] + ) + expected_policy_metadata = build_groot_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + groot_dimensions, + args.text.name, + args.mmproj.name, + ) + elif framework == "pi": + expected_policy_metadata = build_pi_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + dict(PI_SUPPORTED_DIMENSIONS), + args.text.name, + args.mmproj.name, + ) + else: + expected_policy_metadata = build_pi_v3_metadata( + args.policy_dir, + args.hf_dir, + variant, + surgery_manifest, + dict(PI_V3_SUPPORTED_DIMENSIONS), + args.text.name, + args.mmproj.name, + ) + + text_reader = gguf.GGUFReader(args.text) + mmproj_reader = gguf.GGUFReader(args.mmproj) + policy_reader = gguf.GGUFReader(args.policy) + text_validation = validate_text( + text_reader, + bundle_uuid, + args.text_dtype, + args.hf_dir, + require_oft_action_token=framework == "oft", + backbone=backbone, + ) + mmproj_validation = validate_mmproj( + mmproj_reader, + bundle_uuid, + args.mmproj_dtype, + args.hf_dir, + backbone=backbone, + ) + policy_validation = validate_policy( + policy_reader, + bundle_uuid, + args.policy_dtype, + args.policy_dir, + args.text.name, + args.mmproj.name, + expected_policy_metadata, + framework=framework, + backbone=backbone, + ) + del text_reader, mmproj_reader, policy_reader + + source_tensors = surgery_manifest["tensors"] + role_counts = Counter(record["role"] for record in source_tensors) + expected_role_counts = Counter( + { + "text": int(variant["expected"]["text_tensors"]), + "visual": int(variant["expected"]["visual_tensors"]), + "policy": int(variant["expected"]["policy_tensors"]), + "lm_head": int(variant["expected"]["lm_head_tensors"]), + } + ) + if role_counts != expected_role_counts: + raise StarVLAError(f"unexpected surgery source tensor coverage: {dict(role_counts)}") + manifest = { + "schema_version": 1, + "variant": args.variant, + "model_type": variant["model_type"], + "bundle_uuid": bundle_uuid, + "source": portable_source_record(surgery_manifest["source"], variant), + "source_tensor_roles": dict(sorted(role_counts.items())), + "components": { + "text": component_record(args.text, text_validation), + "mmproj": component_record(args.mmproj, mmproj_validation), + "policy": component_record(args.policy, policy_validation), + }, + } + atomic_write_json(args.output, manifest, overwrite=False) + print(f"conversion manifest: {args.output}") + return 0 + except (StarVLAError, OSError, ValueError, TypeError, KeyError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main())