feat: add restricted torch legacy checkpoint loading (#1443)

This commit is contained in:
leejet authored and GitHub committed 2026-04-19 23:09:43 +08:00
1 parent 66143340b6
commit 0a7ae07f94
14 files changed
+1647 -480

No files matched your search

+3 -16
View File
@@ -6,6 +6,7 @@
#include <string>
#include <vector>
#include "binary_io.h"
#include "json.hpp"
static constexpr size_t ST_HEADER_SIZE_LEN = 8;
@@ -16,20 +17,6 @@ static void set_error(std::string* error, const std::string& message) {
}
}
static uint64_t read_u64(const uint8_t* buffer) {
// little endian
uint64_t value = 0;
value |= static_cast<uint64_t>(buffer[7]) << 56;
value |= static_cast<uint64_t>(buffer[6]) << 48;
value |= static_cast<uint64_t>(buffer[5]) << 40;
value |= static_cast<uint64_t>(buffer[4]) << 32;
value |= static_cast<uint64_t>(buffer[3]) << 24;
value |= static_cast<uint64_t>(buffer[2]) << 16;
value |= static_cast<uint64_t>(buffer[1]) << 8;
value |= static_cast<uint64_t>(buffer[0]);
return value;
}
bool is_safetensors_file(const std::string& file_path) {
std::ifstream file(file_path, std::ios::binary);
if (!file.is_open()) {
@@ -52,7 +39,7 @@ bool is_safetensors_file(const std::string& file_path) {
return false;
}
size_t header_size_ = read_u64(header_size_buf);
size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_ || header_size_ <= 2) {
return false;
}
@@ -123,7 +110,7 @@ bool read_safetensors_file(const std::string& file_path,
return false;
}
size_t header_size_ = read_u64(header_size_buf);
size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_) {
set_error(error, "invalid safetensor file '" + file_path + "'");
return false;