mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +00:00
206 lines
8.5 KiB
C++
206 lines
8.5 KiB
C++
#include "model_manager.h"
|
|
|
|
#include <algorithm>
|
|
|
|
#include "core/util.h"
|
|
|
|
static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
|
|
return a.file_id == b.file_id && a.file_revision == b.file_revision &&
|
|
a.file_index == b.file_index && a.offset == b.offset && a.index_in_zip == b.index_in_zip &&
|
|
a.storage_key == b.storage_key && a.type == b.type && a.expected_type == b.expected_type &&
|
|
a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
|
|
a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
|
|
a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
|
|
a.is_int8_tensorwise == b.is_int8_tensorwise && a.w4_convrot_kind == b.w4_convrot_kind && a.int8_convrot == b.int8_convrot &&
|
|
a.int8_convrot_group_size == b.int8_convrot_group_size;
|
|
}
|
|
|
|
std::vector<TensorStorage> ModelManager::find_tensor_sources(const TensorState& state, const String2TensorStorage& sources) {
|
|
auto first = sources.find(state.name);
|
|
if (first == sources.end()) {
|
|
return {};
|
|
}
|
|
std::vector<TensorStorage> result{first->second};
|
|
if (state.component == ModelComponent::LoRA ||
|
|
std::equal(first->second.ne, first->second.ne + GGML_MAX_DIMS, state.tensor->ne)) {
|
|
return result;
|
|
}
|
|
for (size_t i = 1;; ++i) {
|
|
auto part = sources.find(state.name + "." + std::to_string(i));
|
|
if (part == sources.end()) {
|
|
break;
|
|
}
|
|
result.push_back(part->second);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
|
|
auto affected = states;
|
|
for (const auto& block : params_storage_blocks_) {
|
|
if (std::any_of(block->states.begin(), block->states.end(), [&](TensorState* state) { return states.count(state) != 0; })) {
|
|
affected.insert(block->states.begin(), block->states.end());
|
|
}
|
|
}
|
|
for (auto it = prefetch_blocks_.begin(); it != prefetch_blocks_.end();) {
|
|
if (std::any_of(it->second->states.begin(), it->second->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
|
|
free_prefetch_block(*it->second);
|
|
it = prefetch_blocks_.erase(it);
|
|
} else {
|
|
++it;
|
|
}
|
|
}
|
|
for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
|
|
if (std::any_of((*it)->staged_tensors.begin(), (*it)->staged_tensors.end(), [&](const auto& entry) { return affected.count(entry.first) != 0; })) {
|
|
ggml_backend_synchronize((*it)->compute_backend);
|
|
free_compute_staging_block(**it);
|
|
it = compute_staging_blocks_.erase(it);
|
|
} else {
|
|
++it;
|
|
}
|
|
}
|
|
for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
|
|
if (std::any_of((*it)->states.begin(), (*it)->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
|
|
free_params_storage_block(**it);
|
|
it = params_storage_blocks_.erase(it);
|
|
} else {
|
|
++it;
|
|
}
|
|
}
|
|
for (auto* state : affected) {
|
|
state->metadata_validated = false;
|
|
state->applied_lora_epoch = UINT64_MAX;
|
|
}
|
|
}
|
|
|
|
bool ModelManager::set_loader(ModelLoader loader) {
|
|
if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
|
|
return state->pin_count != 0;
|
|
})) {
|
|
LOG_ERROR("cannot update model sources during execution");
|
|
return false;
|
|
}
|
|
std::map<std::pair<ModelLoader::FileId, SDVersion>, String2TensorStorage> scoped;
|
|
auto sources_for = [&](const TensorState& state) -> const String2TensorStorage& {
|
|
if (state.source_file == 0)
|
|
return loader.get_tensor_storage_map();
|
|
auto key = std::make_pair(state.source_file, state.source_version);
|
|
auto found = scoped.find(key);
|
|
if (found == scoped.end())
|
|
found = scoped.emplace(key, model_sources_.count(key.first) != 0
|
|
? source_loader(key.first).file_tensors(key.first, key.second)
|
|
: loader.file_tensors(key.first, key.second))
|
|
.first;
|
|
return found->second;
|
|
};
|
|
bool lora_changed = false;
|
|
for (const auto& spec : loras_) {
|
|
lora_changed |= loader.file_revision(spec.file_id) != spec.file_revision;
|
|
}
|
|
std::unordered_set<TensorState*> changed;
|
|
for (const auto& state : tensor_states_) {
|
|
const auto sources = find_tensor_sources(*state, sources_for(*state));
|
|
if (sources.size() != state->sources.size() ||
|
|
!std::equal(sources.begin(), sources.end(), state->sources.begin(), same_tensor_source) ||
|
|
(lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
|
|
changed.insert(state.get());
|
|
}
|
|
}
|
|
invalidate_sources(changed);
|
|
for (auto* state : changed) {
|
|
state->sources = find_tensor_sources(*state, sources_for(*state));
|
|
}
|
|
if (lora_changed) {
|
|
++current_lora_epoch_;
|
|
for (auto& spec : loras_)
|
|
spec.file_revision = loader.file_revision(spec.file_id);
|
|
}
|
|
model_loader_ = std::move(loader);
|
|
model_loader_.set_n_threads(n_threads_);
|
|
return true;
|
|
}
|
|
|
|
bool ModelManager::add_file(const std::string& path, const std::string& prefix, ModelLoader::FileId* id, bool force) {
|
|
ModelLoader candidate = model_loader_;
|
|
ModelLoader::FileId added_id;
|
|
if (!candidate.add_file(path, prefix, &added_id, force) || !set_loader(std::move(candidate))) {
|
|
return false;
|
|
}
|
|
if (id != nullptr) {
|
|
*id = added_id;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool ModelManager::del_file(ModelLoader::FileId id) {
|
|
ModelLoader candidate = model_loader_;
|
|
return candidate.del_file(id) && set_loader(std::move(candidate));
|
|
}
|
|
|
|
bool ModelManager::refresh_files() {
|
|
ModelLoader candidate = model_loader_;
|
|
return candidate.refresh_files() && set_loader(std::move(candidate));
|
|
}
|
|
|
|
ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const {
|
|
ModelLoader::FileVersions versions;
|
|
const auto& sources = loader.get_tensor_storage_map();
|
|
for (const auto& state : tensor_states_) {
|
|
if (components.count(state->component) == 0) {
|
|
continue;
|
|
}
|
|
if (state->source_file != 0) {
|
|
versions[state->source_file] = model_sources_.count(state->source_file) != 0
|
|
? source_loader(state->source_file).file_revision(state->source_file)
|
|
: loader.file_revision(state->source_file);
|
|
continue;
|
|
}
|
|
for (const auto& source : find_tensor_sources(*state, sources)) {
|
|
versions[source.file_id] = source.file_revision;
|
|
}
|
|
}
|
|
return versions;
|
|
}
|
|
|
|
size_t ModelManager::registered_params_size(const std::set<ModelComponent>& components) const {
|
|
size_t bytes = 0;
|
|
std::unordered_set<const ggml_tensor*> seen;
|
|
for (const auto& state : tensor_states_) {
|
|
if (components.count(state->component) != 0 && state->tensor != nullptr && seen.insert(state->tensor).second) {
|
|
bytes += ggml_nbytes(state->tensor);
|
|
}
|
|
}
|
|
return bytes;
|
|
}
|
|
|
|
const ModelLoader& ModelManager::source_loader(ModelLoader::FileId file) const {
|
|
auto found = model_sources_.find(file);
|
|
return found == model_sources_.end() ? model_loader_ : found->second;
|
|
}
|
|
|
|
ModelLoader::FileId ModelManager::register_model_source(ModelLoader loader) {
|
|
const auto versions = loader.file_versions();
|
|
if (versions.size() != 1) {
|
|
LOG_ERROR("a model source must contain exactly one root model file");
|
|
return 0;
|
|
}
|
|
const auto file = versions.begin()->first;
|
|
if (model_loader_.file_revision(file) != 0 || model_sources_.count(file) != 0) {
|
|
LOG_ERROR("model source is already registered");
|
|
return 0;
|
|
}
|
|
loader.set_n_threads(n_threads_);
|
|
model_sources_.emplace(file, std::move(loader));
|
|
return file;
|
|
}
|
|
|
|
bool ModelManager::unregister_model_source(ModelLoader::FileId file) {
|
|
if (std::any_of(tensor_states_.begin(), tensor_states_.end(), [file](const auto& state) {
|
|
return state->source_file == file;
|
|
})) {
|
|
LOG_ERROR("cannot unregister a model source while its weights are registered");
|
|
return false;
|
|
}
|
|
return model_sources_.erase(file) != 0;
|
|
}
|