mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-26 16:07:50 -05:00
feat: support safetensors index loading (#1769)
This commit is contained in:
@@ -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<std::string> 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) {
|
||||
|
||||
Reference in New Issue
Block a user