upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-28 17:08:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix: prevent torch checkpoint offset overflow (#1832)

Browse Source
This commit is contained in:
leejet
2026-07-29 23:16:29 +08:00
committed by GitHub
parent 9cfe2af8f9
commit e92e86fb11
6 changed files with 121 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/model_loader.cpp
+39 -9
View File
@@ -1077,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::vector<uint8_t> read_buffer;
std::vector<uint8_t> convert_buffer;
std::vector<uint8_t> zip_entry_buffer;
while (true) {
int64_t t0, t1;
@@ -1115,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
auto read_data = [&](char* buf, size_t n) {
auto read_data = [&](char* buf, size_t n) -> bool {
if (zip != nullptr) {
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
return false;
}
size_t entry_size = zip_entry_size(zip);
if (tensor_storage.offset > entry_size) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
if (n > entry_size - tensor_offset) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
if (entry_size != n) {
int64_t t_memcpy_start;
read_buffer.resize(entry_size);
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
zip_entry_buffer.resize(entry_size);
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
t_memcpy_start = ggml_time_ms();
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
} else {
zip_entry_noallocread(zip, (void*)buf, n);
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
}
zip_entry_close(zip);
} else if (mmapped) {
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
} else {
file.seekg(tensor_storage.offset);
file.read(buf, n);
if (!file) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
}
return true;
};
char* read_buf = nullptr;
@@ -1176,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}
t0 = ggml_time_ms();
read_data(read_buf, nbytes_to_read);
if (!read_data(read_buf, nbytes_to_read)) {
failed = true;
break;
}
t1 = ggml_time_ms();
read_time_ms.fetch_add(t1 - t0);