DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-05-08 16:28:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: harden safetensors and gguf loading code (#1404)

Browse Source 
Co-authored-by: professor-moody <keys@nimbus.lan>
This commit is contained in:
Wagner Bruna 2026-04-11 06:19:57 -03:00 committed by GitHub
parent be9f51b25c
commit 118489eb5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/gguf_reader.hpp
View File

@ -59,6 +59,9 @@ private:
if (!safe_read(fin, key_len)) if (!safe_read(fin, key_len))
return false; return false;
if (key_len > 4096)
return false;
std::string key(key_len, '\0'); std::string key(key_len, '\0');
if (!safe_read(fin, (char*)key.data(), key_len)) if (!safe_read(fin, (char*)key.data(), key_len))
return false; return false;

29
src/model.cpp
View File

@ -315,8 +315,9 @@ bool is_safetensors_file(const std::string& file_path) {
if (!file) { if (!file) {
return false; return false;
} }
nlohmann::json header_ = nlohmann::json::parse(header_buf.data()); try {
if (header_.is_discarded()) { nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
} catch (const std::exception&) {
return false; return false;
} }
return true; return true;
@ -511,7 +512,14 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
return false; return false;
} }
nlohmann::json header_ = nlohmann::json::parse(header_buf.data()); nlohmann::json header_;
try {
header_ = nlohmann::json::parse(header_buf.data());
} catch (const std::exception&) {
LOG_ERROR("parsing safetensors header failed", file_path.c_str());
file_paths_.pop_back();
return false;
}
for (auto& item : header_.items()) { for (auto& item : header_.items()) {
std::string name = item.key(); std::string name = item.key();
@ -575,24 +583,29 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
size_t tensor_data_size = end - begin; size_t tensor_data_size = end - begin;
bool tensor_size_ok;
if (dtype == "F8_E4M3") { if (dtype == "F8_E4M3") {
tensor_storage.is_f8_e4m3 = true; tensor_storage.is_f8_e4m3 = true;
// f8 -> f16 // f8 -> f16
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2); tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
} else if (dtype == "F8_E5M2") { } else if (dtype == "F8_E5M2") {
tensor_storage.is_f8_e5m2 = true; tensor_storage.is_f8_e5m2 = true;
// f8 -> f16 // f8 -> f16
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2); tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
} else if (dtype == "F64") { } else if (dtype == "F64") {
tensor_storage.is_f64 = true; tensor_storage.is_f64 = true;
// f64 -> f32 // f64 -> f32
GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size); tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
} else if (dtype == "I64") { } else if (dtype == "I64") {
tensor_storage.is_i64 = true; tensor_storage.is_i64 = true;
// i64 -> i32 // i64 -> i32
GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size); tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
} else { } else {
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size); tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
}
if (!tensor_size_ok) {
LOG_ERROR("size mismatch for tensor '%s' (%s)\n", name.c_str(), dtype.c_str());
return false;
} }
add_tensor_storage(tensor_storage); add_tensor_storage(tensor_storage);