mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-14 12:56:33 +00:00
feat: add restricted torch legacy checkpoint loading (#1443)
This commit is contained in:
1 parent
66143340b6
commit
0a7ae07f94
14 files changed
+1647
-480
No files matched your search
@@ -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;
|
||||
|
||||
Reference in new issue
Block a user