From ab0b73d5263a20898e43c11b142778b837556790 Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 10 Jul 2026 21:18:25 +0800 Subject: [PATCH] feat: support safetensors index loading --- src/model_io/safetensors_io.cpp | 58 +++++++++++++++++++++++++++++++++ src/model_io/safetensors_io.h | 3 ++ src/model_loader.cpp | 22 +++++++++++++ src/model_loader.h | 1 + 4 files changed, 84 insertions(+) diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 197dc5f54..223cc9670 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -3,15 +3,19 @@ #include #include #include +#include #include #include #include +#include #include #include "binary_io.h" #include "core/util.h" #include "json.hpp" +namespace fs = std::filesystem; + static constexpr size_t ST_HEADER_SIZE_LEN = 8; static void set_error(std::string* error, const std::string& message) { @@ -20,6 +24,14 @@ static void set_error(std::string* error, const std::string& message) { } } +static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) { + fs::path shard_fs_path(shard_path); + if (shard_fs_path.is_absolute()) { + return shard_fs_path.lexically_normal().string(); + } + return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string(); +} + bool is_safetensors_file(const std::string& file_path) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { @@ -230,6 +242,52 @@ bool read_safetensors_file(const std::string& file_path, return true; } +bool read_safetensors_index_file(const std::string& file_path, + std::vector& shard_paths, + std::string* error) { + shard_paths.clear(); + + std::ifstream file(file_path); + if (!file.is_open()) { + set_error(error, "failed to open '" + file_path + "'"); + return false; + } + + nlohmann::json index; + try { + index = nlohmann::json::parse(file); + } catch (const std::exception&) { + set_error(error, "parsing safetensors index failed: '" + file_path + "'"); + return false; + } + + if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) { + set_error(error, "invalid safetensors index '" + file_path + "'"); + return false; + } + + std::unordered_set seen_shard_paths; + for (const auto& item : index["weight_map"].items()) { + if (!item.value().is_string()) { + set_error(error, "invalid shard path for tensor '" + item.key() + "'"); + return false; + } + + std::string shard_path = resolve_index_shard_path(file_path, + item.value().get()); + if (seen_shard_paths.insert(shard_path).second) { + shard_paths.push_back(std::move(shard_path)); + } + } + + if (shard_paths.empty()) { + set_error(error, "safetensors index has no tensors: '" + file_path + "'"); + return false; + } + + return true; +} + static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) { switch (type) { case GGML_TYPE_F16: diff --git a/src/model_io/safetensors_io.h b/src/model_io/safetensors_io.h index b4938ee18..b18f0ae75 100644 --- a/src/model_io/safetensors_io.h +++ b/src/model_io/safetensors_io.h @@ -11,6 +11,9 @@ bool is_safetensors_file(const std::string& file_path); bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, std::string* error = nullptr); +bool read_safetensors_index_file(const std::string& file_path, + std::vector& shard_paths, + std::string* error = nullptr); bool write_safetensors_file(const std::string& file_path, const std::vector& tensors, std::string* error = nullptr); diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 51091d98f..0dc7774af 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -235,6 +235,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string } else if (is_gguf_file(file_path)) { LOG_INFO("load %s using gguf format", file_path.c_str()); return init_from_gguf_file(file_path, prefix); + } else if (ends_with(file_path, ".json")) { + LOG_INFO("load %s using safetensors index format", file_path.c_str()); + return init_from_safetensors_index_file(file_path, prefix); } else if (is_safetensors_file(file_path)) { LOG_INFO("load %s using safetensors format", file_path.c_str()); return init_from_safetensors_file(file_path, prefix); @@ -339,6 +342,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const return true; } +bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) { + LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); + + std::vector shard_paths; + std::string error; + if (!read_safetensors_index_file(file_path, shard_paths, &error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + + for (const std::string& shard_path : shard_paths) { + if (!init_from_file(shard_path, prefix)) { + return false; + } + } + + return true; +} + /*================================================= TorchLegacyModelLoader ==================================================*/ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) { diff --git a/src/model_loader.h b/src/model_loader.h index 6c19c53c5..d311b7030 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -43,6 +43,7 @@ class ModelLoader { bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");