stable-diffusion.cpp/src/model_manager_files.cpp

#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.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 (!workspace_reclaimers_.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, 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] = 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;
}