DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-24 20:20:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: unify model source and weight lifecycle management (#1956)

Browse Source
This commit is contained in:
leejet 2026-09-10 23:53:48 +08:00 committed by GitHub
parent d04e8950c1
commit 6b47fec013
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
36 changed files with 2548 additions and 1462 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
src/core/ggml_runner.cpp
View File

@ -66,15 +66,9 @@ void GGMLRunner::rebuild_params_tensor_set() {
}
ggml_tensor* GGMLRunner::canonical_param_tensor(ggml_tensor* tensor) {
if (tensor == nullptr) {
return nullptr;
}
if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
return tensor;
}
if (tensor->view_src != nullptr &&
params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
return tensor->view_src;
for (auto* current = tensor; current != nullptr; current = current->view_src) {
if (params_tensor_set_.count(current) != 0)
return current;
}
return nullptr;
}
@ -483,9 +477,10 @@ void GGMLRunner::runner_end() {
if (auto manager = residency_manager.lock()) {
manager->clear_prefetched_params(reinterpret_cast<uintptr_t>(this));
std::vector<ggml_tensor*> tensors;
for (auto tensor = ggml_get_first_tensor(params_ctx); tensor != nullptr;
tensor = ggml_get_next_tensor(params_ctx, tensor)) {
tensors.push_back(tensor);
for (auto tensor : params_tensor_set_) {
auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
if (parameter != nullptr)
tensors.push_back(parameter);
}
manager->evict_compute_backend_params(tensors);
manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(this));
@ -620,7 +615,15 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
if (!prepare_compute_graph(get_graph, &graph)) {
return std::nullopt;
}
params_tensor_set_dirty_ = true;
rebuild_params_tensor_set();
if (auto manager = residency_manager.lock()) {
for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) {
auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i));
if (parameter != nullptr)
params_tensor_set_.insert(parameter);
}
}
auto output = execute_graph(graph, n_threads, no_return, read_outputs);
success = output.has_value();
if (success) {

18
src/core/ordered_map.hpp
View File

@ -38,9 +38,19 @@ public:
insert(kv);
}
OrderedMap(const OrderedMap&) = default;
OrderedMap(OrderedMap&&) noexcept = default;
OrderedMap& operator=(const OrderedMap&) = default;
OrderedMap(const OrderedMap& other) {
for (const auto& value : other) {
insert(value);
}
}
OrderedMap(OrderedMap&&) noexcept = default;
OrderedMap& operator=(const OrderedMap& other) {
if (this != &other) {
OrderedMap copy(other);
swap(copy);
}
return *this;
}
OrderedMap& operator=(OrderedMap&&) noexcept = default;
// --- element access ---
@ -174,4 +184,4 @@ public:
}
};
#endif // __SD_CORE_ORDERED_MAP_HPP__
#endif // __SD_CORE_ORDERED_MAP_HPP__

5
src/detailer.cpp
View File

@ -676,7 +676,7 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
model_manager = std::make_shared<ModelManager>();
model_manager->set_n_threads(n_threads);
model_manager->set_enable_mmap(false);
ModelLoader& loader = model_manager->loader();
ModelLoader loader;
if (!loader.init_from_file(detector_path)) {
LOG_ERROR("failed to load ADetailer detector: '%s'", detector_path.c_str());
return false;
@ -696,7 +696,8 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) {
std::map<std::string, ggml_tensor*> tensors;
detector->get_param_tensors(tensors);
if (!model_manager->register_param_tensors("YOLOv8",
if (!model_manager->set_loader(loader) ||
!model_manager->register_param_tensors(ModelComponent::Detector,
std::move(tensors),
backend_manager.params_backend_is_disk(SDBackendModule::DETECTOR)
? ModelManager::ResidencyMode::Disk

7
src/device_residency_manager.h
View File

@ -48,9 +48,10 @@ struct DeviceResidencyManager {
const std::vector<ggml_tensor*>& required_params) const = 0;
virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) = 0;
virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const { return nullptr; }
virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual WeightResidencyInfo inspect_compute_backend_params(
const std::vector<ggml_tensor*>& tensors) const = 0;
virtual void update_runtime_residency(uintptr_t owner_id,

5
src/extensions/generation_extension.h
View File

@ -19,7 +19,7 @@ struct GenerationExtensionInitContext {
const sd_ctx_params_t* params;
SDVersion version;
const String2TensorStorage& tensor_storage_map;
ModelLoader& model_loader;
bool photomaker_source_available;
std::shared_ptr<ModelManager> model_manager;
int n_threads;
std::function<bool(SDBackendModule)> ensure_backend_pair;
@ -39,7 +39,8 @@ struct GenerationExtensionConditionContext {
struct GenerationExtension {
virtual ~GenerationExtension() = default;
virtual const char* name() const = 0;
virtual ModelComponent component() const = 0;
const char* name() const { return model_component_name(component()); }
virtual bool is_enabled() const {
return false;
}

14
src/extensions/photomaker_extension.cpp
View File

@ -109,8 +109,8 @@ struct PhotoMakerExtension : public GenerationExtension {
SDCondition id_condition;
int start_merge_step = -1;
const char* name() const override {
return "photomaker";
ModelComponent component() const override {
return ModelComponent::PhotoMaker;
}
bool is_enabled() const override {
@ -119,7 +119,7 @@ struct PhotoMakerExtension : public GenerationExtension {
bool init(const GenerationExtensionInitContext& ctx) override {
model_path = SAFE_STR(ctx.params->photo_maker_path);
if (model_path.empty()) {
if (model_path.empty() || !ctx.photomaker_source_available) {
return true;
}
@ -128,13 +128,7 @@ struct PhotoMakerExtension : public GenerationExtension {
}
PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1;
LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", model_path.c_str());
if (!ctx.model_loader.init_from_file_and_convert_name(model_path, "pmid.")) {
LOG_WARN("loading stacked ID embedding from '%s' failed", model_path.c_str());
return true;
}
pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
ctx.tensor_storage_map,
"pmid",
ctx.version,

4
src/extensions/pulid_extension.cpp
View File

@ -79,8 +79,8 @@ struct PuLIDExtension : public GenerationExtension {
sd::Tensor<float> id_embedding;
float id_weight = 1.0f;
const char* name() const override {
return "pulid";
ModelComponent component() const override {
return ModelComponent::PuLID;
}
bool is_enabled() const override {

183
src/model/adapter/lora.hpp
View File

@ -23,25 +23,31 @@ struct LoraModel : public GGMLRunner {
std::set<std::string> skipped_incompatible_lora_tensors;
std::set<std::string> warned_incompatible_model_tensors;
std::string file_path;
std::shared_ptr<ModelManager> model_manager;
ggml_backend_t params_backend = nullptr;
bool load_failed = false;
bool applied = false;
bool tensor_preprocessed = false;
ggml_backend_t params_backend = nullptr;
bool load_failed = false;
bool applied = false;
bool tensor_preprocessed = false;
ModelLoader::FileId source_file = 0;
SDVersion source_version = VERSION_COUNT;
ModelManager::ResidencyMode residency_mode = ModelManager::ResidencyMode::ParamBackend;
bool params_follow_compute = false;
std::vector<ggml_tensor*> registered_params;
std::map<ggml_tensor*, float> scalar_values;
typedef std::function<bool(const std::string&)> filter_t;
LoraModel(const std::string& lora_id,
ggml_backend_t backend,
ggml_backend_t params_backend_,
const std::string& file_path = "",
std::string prefix = "",
SDVersion version = VERSION_COUNT,
std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>())
: GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
prefix = "lora." + prefix;
if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) {
load_failed = true;
LoraModel(const std::string& id, ggml_backend_t backend, ggml_backend_t params, std::shared_ptr<ModelManager> manager, ModelLoader::FileId file, SDVersion version, ModelManager::ResidencyMode mode = ModelManager::ResidencyMode::ParamBackend, bool follow_compute = false)
: GGMLRunner(backend, manager), lora_id(id), params_backend(params), source_file(file), source_version(version), residency_mode(mode), params_follow_compute(follow_compute) {
load_failed = source_file == 0 || manager == nullptr || manager->loader().file_revision(source_file) == 0;
if (!load_failed) {
file_path = manager->loader().file_path(source_file);
}
}
~LoraModel() override {
runner_end();
if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
GGML_ASSERT(manager->unregister_param_tensors(registered_params));
}
}
@ -49,95 +55,65 @@ struct LoraModel : public GGMLRunner {
return "lora";
}
bool load_from_file(int n_threads, filter_t filter = nullptr) {
LOG_INFO("loading LoRA from '%s'", file_path.c_str());
if (load_failed) {
LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
bool init_params(int n_threads, filter_t filter = nullptr) {
auto model_manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
if (model_manager == nullptr)
return false;
}
std::unordered_map<std::string, TensorStorage> tensors_to_create;
std::mutex lora_mutex;
bool dry_run = true;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
if (dry_run) {
const std::string& name = tensor_storage.name;
if (filter && !filter(name)) {
return true;
}
{
std::lock_guard<std::mutex> lock(lora_mutex);
tensors_to_create[name] = tensor_storage;
}
} else {
const std::string& name = tensor_storage.name;
auto iter = lora_tensors.find(name);
if (iter != lora_tensors.end()) {
*dst_tensor = iter->second;
}
}
return true;
};
if (model_manager != nullptr) {
model_manager->set_n_threads(n_threads);
}
ModelLoader& model_loader = model_manager->loader();
model_loader.load_tensors(on_new_tensor_cb);
if (tensors_to_create.empty()) {
return true;
}
for (const auto& pair : tensors_to_create) {
const auto& name = pair.first;
const auto& ts = pair.second;
ggml_tensor* real = ggml_new_tensor(params_ctx,
ts.type,
ts.n_dims,
ts.ne);
lora_tensors[name] = real;
}
if (load_failed || !registered_params.empty())
return false;
model_manager->set_n_threads(n_threads);
const auto sources = model_manager->loader().file_tensors(source_file, source_version);
std::map<std::string, ggml_tensor*> tensors;
for (const auto& pair : lora_tensors) {
tensors[pair.first] = pair.second;
std::map<std::string, ggml_tensor*> scalars;
std::set<std::string> scalar_names;
for (const auto& [name, source] : sources) {
if (is_unused_tensor(name) || (filter && !filter(name)))
continue;
const bool scalar = source.nelements() == 1 && (ends_with(name, ".alpha") || ends_with(name, ".scale"));
auto* tensor = ggml_new_tensor(params_ctx, scalar ? GGML_TYPE_F32 : source.type, source.n_dims, source.ne);
lora_tensors[name] = tensor;
if (scalar) {
tensor->data = &scalar_values[tensor];
scalars[name] = tensor;
scalar_names.insert(name);
} else {
tensors[name] = tensor;
}
}
if (model_manager == nullptr ||
!model_manager->register_param_tensors("LoRA",
std::move(tensors),
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
LOG_ERROR("lora model manager registration failed");
// These values are consumed while constructing the graph, before weight preparation.
if (!scalars.empty()) {
auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
auto found = scalars.find(source.name);
*dst = found == scalars.end() ? nullptr : found->second;
return true;
};
if (!model_manager->loader().load_file_tensors(source_file, source_version, callback, scalar_names))
return false;
}
if (!model_manager->register_param_tensors(ModelComponent::LoRA, tensors, residency_mode,
runtime_backend, params_backend, nullptr, false, params_follow_compute,
nullptr, source_file, source_version))
return false;
}
std::vector<ggml_tensor*> lora_params;
lora_params.reserve(lora_tensors.size());
for (const auto& pair : lora_tensors) {
lora_params.push_back(pair.second);
}
if (!model_manager->prepare_params(lora_params)) {
LOG_ERROR("lora model manager prepare params failed");
return false;
}
for (const auto& entry : tensors)
registered_params.push_back(entry.second);
return model_manager->validate_registered_tensors();
}
LOG_VERBOSE("finished loaded lora");
return true;
float scalar_value(ggml_tensor* tensor) const {
auto found = scalar_values.find(tensor);
return found != scalar_values.end() ? found->second : ggml_ext_backend_tensor_get_f32(tensor);
}
void release_loaded_tensors() {
runner_end();
model_manager.reset();
if (auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock())) {
GGML_ASSERT(manager->unregister_param_tensors(registered_params));
}
registered_params.clear();
free_params_ctx();
alloc_params_ctx();
model_manager = std::make_shared<ModelManager>();
residency_manager = model_manager;
lora_tensors.clear();
scalar_values.clear();
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
skipped_incompatible_lora_tensors.clear();
@ -241,12 +217,12 @@ struct LoraModel : public GGMLRunner {
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = scalar_value(iter->second);
applied_lora_tensors.insert(scale_name);
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
float alpha = scalar_value(iter->second);
scale_value = alpha / rank;
// LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
applied_lora_tensors.insert(alpha_name);
@ -395,7 +371,7 @@ struct LoraModel : public GGMLRunner {
int64_t rank = hada_1_down->ne[ggml_n_dims(hada_1_down) - 1];
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
float alpha = scalar_value(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
@ -508,7 +484,7 @@ struct LoraModel : public GGMLRunner {
float scale_value = 1.0f;
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
float alpha = scalar_value(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
@ -669,7 +645,7 @@ struct LoraModel : public GGMLRunner {
float scale_value = 1.0f;
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
float alpha = scalar_value(iter->second);
scale_value = alpha / rank;
}
@ -796,12 +772,12 @@ struct LoraModel : public GGMLRunner {
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = scalar_value(iter->second);
scale_tensor_name = scale_name;
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
float alpha = scalar_value(iter->second);
scale_value = alpha / rank;
scale_tensor_name = alpha_name;
// LOG_VERBOSE("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
@ -949,7 +925,7 @@ struct LoraModel : public GGMLRunner {
return gf;
}
void apply(std::map<std::string, ggml_tensor*> model_tensors,
bool apply(std::map<std::string, ggml_tensor*> model_tensors,
const std::set<std::string>& model_tensor_names,
SDVersion version,
int n_threads,
@ -970,10 +946,11 @@ struct LoraModel : public GGMLRunner {
stat(!warn_unused);
original_tensor_to_final_tensor.clear();
runner_end();
return result.has_value();
}
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
bool apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
return apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
}
void stat(bool at_runntime = false) {

92
src/model/adapter/pmid.hpp
View File

@ -6,10 +6,8 @@
#include "core/util.h"
#include "model/common/ggml_block.hpp"
#include "model/adapter/lora.hpp"
#include "model/common/block.hpp"
#include "model/te/clip.hpp"
#include "model_loader.h"
struct FuseBlock : public GGMLBlock {
// network hparams
@ -565,94 +563,4 @@ public:
}
};
struct PhotoMakerIDEmbed : public GGMLRunner {
std::map<std::string, ggml_tensor*> tensors;
std::string file_path;
std::shared_ptr<ModelManager> model_manager;
ggml_backend_t params_backend = nullptr;
bool load_failed = false;
bool applied = false;
PhotoMakerIDEmbed(ggml_backend_t backend,
ggml_backend_t params_backend_,
std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>(),
const std::string& file_path = "",
const std::string& prefix = "")
: GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) {
load_failed = true;
}
}
std::string get_desc() {
return "id_embeds";
}
bool load_from_file(bool filter_tensor, int n_threads) {
LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
if (load_failed) {
LOG_ERROR("init photomaker id embed from file failed: '%s'", file_path.c_str());
return false;
}
bool dry_run = true;
std::mutex tensor_mutex;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
const std::string& name = tensor_storage.name;
if (filter_tensor && !contains(name, "pmid.id_embeds")) {
// LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
return true;
}
if (dry_run) {
std::lock_guard<std::mutex> lock(tensor_mutex);
ggml_tensor* real = ggml_new_tensor(params_ctx,
tensor_storage.type,
tensor_storage.n_dims,
tensor_storage.ne);
tensors[name] = real;
} else {
auto real = tensors[name];
*dst_tensor = real;
}
return true;
};
model_manager->set_n_threads(n_threads);
ModelLoader& model_loader = model_manager->loader();
model_loader.load_tensors(on_new_tensor_cb);
if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
tensors,
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
LOG_ERROR("PhotoMaker ID embeds model manager registration failed");
return false;
}
std::vector<ggml_tensor*> id_embed_params;
id_embed_params.reserve(tensors.size());
for (const auto& pair : tensors) {
id_embed_params.push_back(pair.second);
}
if (!model_manager->prepare_params(id_embed_params)) {
LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed");
return false;
}
LOG_VERBOSE("finished loading PhotoMaker ID Embeds ");
return true;
}
ggml_tensor* get() {
std::map<std::string, ggml_tensor*>::iterator pos;
pos = tensors.find("pmid.id_embeds");
if (pos != tensors.end())
return pos->second;
return nullptr;
}
};
#endif // __SD_MODEL_ADAPTER_PMID_HPP__

40
src/model/diffusion/control.hpp
View File

@ -2,8 +2,6 @@
#define __SD_MODEL_DIFFUSION_CONTROL_HPP__
#include "model/common/block.hpp"
#include "model_loader.h"
#include "model_manager.h"
// Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
#define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
@ -317,20 +315,17 @@ struct ControlNet : public GGMLRunner {
ggml_tensor* guided_hint_output_ggml = nullptr;
std::vector<sd::Tensor<float>> controls;
bool guided_hint_cached = false;
std::shared_ptr<ModelManager> owned_model_manager;
ggml_backend_t params_backend = nullptr;
static const char* guided_hint_cache_name() {
return "controlnet.guided_hint";
}
ControlNet(ggml_backend_t backend,
ggml_backend_t params_backend_,
const String2TensorStorage& tensor_storage_map = {},
SDVersion version = VERSION_SD1,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix), params_backend(params_backend_) {
: GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) {
control_net.init(params_ctx, tensor_storage_map, prefix);
}
@ -445,39 +440,6 @@ struct ControlNet : public GGMLRunner {
guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
return controls;
}
bool load_from_file(const std::string& file_path, int n_threads) {
LOG_INFO("loading control net from '%s'", file_path.c_str());
std::map<std::string, ggml_tensor*> tensors;
control_net.get_param_tensors(tensors);
auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
if (manager == nullptr) {
owned_model_manager = std::make_shared<ModelManager>();
residency_manager = owned_model_manager;
manager = owned_model_manager;
}
ModelLoader& model_loader = manager->loader();
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
return false;
}
manager->set_n_threads(n_threads);
if (!manager->register_param_tensors("ControlNet",
std::move(tensors),
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!manager->validate_registered_tensors()) {
LOG_ERROR("register control net tensors with model manager failed");
return false;
}
LOG_INFO("control net model loaded");
return true;
}
};
#endif // __SD_MODEL_DIFFUSION_CONTROL_HPP__

7
src/model/diffusion/flux.hpp
View File

@ -1714,8 +1714,8 @@ namespace Flux {
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_COUNT;
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -1736,7 +1736,8 @@ namespace Flux {
VERSION_FLUX2,
model_manager);
if (!model_manager->register_runner_params("Flux test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*flux,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/diffusion/ltxv.hpp
View File

@ -2087,8 +2087,8 @@ namespace LTXV {
ggml_backend_t backend = sd_backend_cpu_init();
LOG_INFO("loading ltxav from '%s'", model_path.c_str());
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(model_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
return;
@ -2107,7 +2107,8 @@ namespace LTXV {
"model.diffusion_model",
model_manager);
if (!model_manager->register_runner_params("LTXAV test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*ltxav,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

5
src/model/diffusion/mmdit.hpp
View File

@ -1064,13 +1064,14 @@ struct MMDiTRunner : public DiffusionModelRunner {
{
LOG_INFO("loading from '%s'", file_path.c_str());
ModelLoader& model_loader = model_manager->loader();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
}
if (!model_manager->register_runner_params("MMDiT test",
if (!model_manager->set_loader(std::move(model_loader)) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*mmdit,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/diffusion/qwen_image.hpp
View File

@ -773,8 +773,8 @@ namespace Qwen {
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -793,7 +793,8 @@ namespace Qwen {
VERSION_QWEN_IMAGE,
model_manager);
if (!model_manager->register_runner_params("Qwen image test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*qwen_image,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/diffusion/wan.hpp
View File

@ -1020,8 +1020,8 @@ namespace WAN {
ggml_type model_data_type = GGML_TYPE_F16;
LOG_INFO("loading from '%s'", file_path.c_str());
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -1040,7 +1040,8 @@ namespace WAN {
VERSION_WAN2_2_TI2V,
model_manager);
if (!model_manager->register_runner_params("Wan test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*wan,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/diffusion/z_image.hpp
View File

@ -706,8 +706,8 @@ namespace ZImage {
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -728,7 +728,8 @@ namespace ZImage {
VERSION_QWEN_IMAGE,
model_manager);
if (!model_manager->register_runner_params("ZImage test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Diffusion,
*z_image,
"model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/te/llm.hpp
View File

@ -2576,8 +2576,8 @@ namespace LLM {
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_COUNT;
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "text_encoders.llm.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -2601,7 +2601,8 @@ namespace LLM {
true,
model_manager);
if (!model_manager->register_runner_params("LLM test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Conditioner,
*llm,
"text_encoders.llm",
ModelManager::ResidencyMode::ParamBackend,

7
src/model/te/t5.hpp
View File

@ -631,8 +631,8 @@ struct T5Embedder {
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_F16;
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
@ -647,7 +647,8 @@ struct T5Embedder {
std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, tensor_storage_map, "", true, model_manager);
if (!model_manager->register_runner_params("T5 test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::Conditioner,
*t5,
"",
ModelManager::ResidencyMode::ParamBackend,

9
src/model/upscaler/ltx_latent_upscaler.hpp
View File

@ -433,12 +433,15 @@ namespace LTXVUpsampler {
struct LatentUpsamplerRunner : public GGMLRunner {
LatentUpsamplerConfig config;
std::unique_ptr<LatentUpsampler> model;
std::string weight_prefix;
LatentUpsamplerRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map)) {
config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map, prefix)),
weight_prefix(prefix) {
if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
@ -452,7 +455,7 @@ namespace LTXVUpsampler {
}
model = std::make_unique<LatentUpsampler>(config);
model->init(params_ctx, tensor_storage_map, "");
model->init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
@ -461,7 +464,7 @@ namespace LTXVUpsampler {
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
if (model) {
model->get_param_tensors(tensors);
model->get_param_tensors(tensors, weight_prefix);
}
}

7
src/model/vae/ltx_audio_vae.hpp
View File

@ -1078,8 +1078,8 @@ namespace LTXV {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file(model_path)) {
LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
return;
@ -1091,7 +1091,8 @@ namespace LTXV {
prefix,
model_manager);
if (!model_manager->register_runner_params("LTX audio VAE test",
if (!model_manager->set_loader(std::move(model_loader)) ||
!model_manager->register_runner_params(ModelComponent::AudioVAE,
*ltx_audio_vae,
ModelManager::ResidencyMode::ParamBackend,
backend,

7
src/model/vae/ltx_vae.hpp
View File

@ -1476,8 +1476,8 @@ struct LTXVideoVAE : public VAE {
ggml_backend_t backend = sd_backend_cpu_init();
LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
auto model_manager = std::make_shared<ModelManager>();
ModelLoader& model_loader = model_manager->loader();
auto model_manager = std::make_shared<ModelManager>();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) {
LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
return;
@ -1491,7 +1491,8 @@ struct LTXVideoVAE : public VAE {
VERSION_LTXAV,
model_manager);
if (!model_manager->register_runner_params("LTX VAE test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::VAE,
*vae,
ModelManager::ResidencyMode::ParamBackend,
backend,

5
src/model/vae/wan_vae.hpp
View File

@ -1494,13 +1494,14 @@ namespace WAN {
{
LOG_INFO("loading from '%s'", file_path.c_str());
ModelLoader& model_loader = model_manager->loader();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
}
if (!model_manager->register_runner_params("Wan VAE test",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_runner_params(ModelComponent::VAE,
*vae,
ModelManager::ResidencyMode::ParamBackend,
backend,

574
src/model_builders.cpp Normal file
View File

@ -0,0 +1,574 @@
#include "model_builders.h"
#include <cstring>
#include <utility>
#include "conditioning/conditioner.hpp"
#include "core/ggml_extend_backend.h"
#include "core/util.h"
#include "extensions/generation_extension.h"
#include "model/adapter/ip_adapter.hpp"
#include "model/diffusion/anima.hpp"
#include "model/diffusion/boogu.hpp"
#include "model/diffusion/control.hpp"
#include "model/diffusion/ernie_image.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/hidream_o1.hpp"
#include "model/diffusion/hunyuan.hpp"
#include "model/diffusion/ideogram4.hpp"
#include "model/diffusion/krea2.hpp"
#include "model/diffusion/lens.hpp"
#include "model/diffusion/lingbot_video.hpp"
#include "model/diffusion/ltxv.hpp"
#include "model/diffusion/mage_flow.hpp"
#include "model/diffusion/minimax_h3.hpp"
#include "model/diffusion/minit2i.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/pid.hpp"
#include "model/diffusion/qwen_image.hpp"
#include "model/diffusion/unet.hpp"
#include "model/diffusion/wan.hpp"
#include "model/diffusion/z_image.hpp"
#include "model/vae/auto_encoder_kl.hpp"
#include "model/vae/hunyuan_vae.hpp"
#include "model/vae/ltx_audio_vae.hpp"
#include "model/vae/ltx_vae.hpp"
#include "model/vae/mage_vae.hpp"
#include "model/vae/minimax_h3_audio_vae.hpp"
#include "model/vae/minimax_h3_vae.hpp"
#include "model/vae/tae.hpp"
#include "model/vae/vae.hpp"
#include "model/vae/wan_vae.hpp"
namespace sd::model_builders {
static bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) {
if (backends.runtime_backend(module) == nullptr) {
LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
return false;
}
if (backends.params_backend(module) == nullptr) {
LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
return false;
}
return true;
}
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
switch (format) {
case SD_VAE_FORMAT_FLUX:
return VERSION_FLUX;
case SD_VAE_FORMAT_SD3:
return VERSION_SD3;
case SD_VAE_FORMAT_FLUX2:
return VERSION_FLUX2;
case SD_VAE_FORMAT_WAN:
return VERSION_WAN2;
case SD_VAE_FORMAT_AUTO:
default:
return fallback;
}
}
bool build_core_runners(const Context& ctx, CoreRunners& runners) {
const auto* sd_ctx_params = &ctx.params;
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
const auto& weight_manager = ctx.weight_manager;
CoreRunners result;
if (!ensure_backend_pair(ctx.backends, SDBackendModule::TE) ||
!ensure_backend_pair(ctx.backends, SDBackendModule::DIFFUSION)) {
return false;
}
if (sd_version_is_sd3(version)) {
result.conditioner = std::make_shared<SD3CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
result.diffusion = std::make_shared<MMDiTRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_pid(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<Pid::PiDRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.net",
weight_manager);
} else if (sd_version_is_ideogram4(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<Ideogram4::Ideogram4Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_krea2(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
weight_manager);
result.diffusion = std::make_shared<Krea2::Krea2Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_flux(version)) {
bool is_chroma = false;
for (auto pair : tensor_storage_map) {
if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
is_chroma = true;
break;
}
}
if (is_chroma) {
result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
false,
1,
false,
weight_manager,
sd_ctx_params->model_args);
} else if (version == VERSION_OVIS_IMAGE) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
} else {
result.conditioner = std::make_shared<FluxCLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
}
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
bool is_chroma = false;
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_ltxav(version)) {
result.conditioner = std::make_shared<LTXAVEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
"text_encoders.llm",
"text_embedding_projection",
weight_manager);
result.diffusion = std::make_shared<LTXV::LTXAVRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_minimax_h3(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
weight_manager);
result.diffusion = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_hunyuan_video(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<Hunyuan::HunyuanVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager);
} else if (sd_version_is_wan(version)) {
result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
true,
0,
true,
weight_manager);
result.diffusion = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager);
if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {
result.high_noise_diffusion = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.high_noise_diffusion_model",
version,
weight_manager);
}
if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" ||
result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" ||
result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") {
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
return false;
}
result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
tensor_storage_map,
weight_manager);
}
} else if (sd_version_is_lingbot_video(version)) {
bool enable_vision = false;
for (const auto& [name, _] : tensor_storage_map) {
if (starts_with(name, "text_encoders.llm.visual.")) {
enable_vision = true;
break;
}
}
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
enable_vision,
weight_manager);
result.diffusion = std::make_shared<LingBotVideo::LingBotVideoRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_qwen_image(version)) {
bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
enable_vision,
weight_manager);
result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_mage_flow(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
weight_manager);
result.diffusion = std::make_shared<MageFlow::MageFlowRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_longcat(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
weight_manager);
result.diffusion = std::make_shared<Flux::FluxRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager,
sd_ctx_params->model_args);
} else if (version == VERSION_HIDREAM_O1) {
result.conditioner = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
result.diffusion = std::make_shared<HiDreamO1::HiDreamO1Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model",
weight_manager);
} else if (sd_version_is_minit2i(version)) {
result.conditioner = std::make_shared<MiniT2IConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
result.diffusion = std::make_shared<MiniT2I::MiniT2IRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.model.net",
weight_manager);
} else if (sd_version_is_anima(version)) {
result.conditioner = std::make_shared<AnimaConditioner>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
weight_manager);
result.diffusion = std::make_shared<Anima::AnimaRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_z_image(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<ZImage::ZImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager);
} else if (sd_version_is_boogu_image(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
true,
weight_manager);
result.diffusion = std::make_shared<Boogu::BooguImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager);
} else if (sd_version_is_ernie_image(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<ErnieImage::ErnieImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_lens(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
weight_manager);
result.diffusion = std::make_shared<Lens::LensRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else { // SD1.x SD2.x SDXL
std::map<std::string, std::string> embbeding_map;
for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path));
}
result.conditioner = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
embbeding_map,
version,
weight_manager);
result.diffusion = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
weight_manager);
if (sd_ctx_params->diffusion_conv_direct) {
LOG_INFO("Using Conv2d direct in the diffusion model");
result.diffusion->set_conv2d_direct_enabled(true);
}
}
if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) {
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
return false;
}
result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
tensor_storage_map,
weight_manager);
}
if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
result.ip_adapter = std::make_shared<IPAdapter::IPAdapterRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"ip_adapter",
weight_manager);
}
runners = std::move(result);
return true;
}
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) {
const auto* sd_ctx_params = &ctx.params;
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
const auto& weight_manager = ctx.weight_manager;
VAERunners result;
if (!ensure_backend_pair(ctx.backends, SDBackendModule::VAE)) {
return false;
}
auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) {
return std::make_shared<TinyVideoAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"decoder",
decode_only,
version,
weight_manager);
} else {
auto model = std::make_shared<TinyImageAutoEncoder>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"decoder.layers",
decode_only,
version,
weight_manager);
return model;
}
};
sd_vae_format_t vae_format = sd_ctx_params->vae_format;
if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
LOG_WARN("invalid VAE format override, using auto");
vae_format = SD_VAE_FORMAT_AUTO;
}
SDVersion vae_version = version;
if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) {
vae_version = sd_vae_format_to_version(vae_format, vae_version);
}
auto create_vae = [&]() -> std::shared_ptr<VAE> {
if (sd_version_is_ltxav(version)) {
return std::make_shared<LTXVideoVAE>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
false,
version,
weight_manager);
} else if (sd_version_is_minimax_h3(version)) {
return std::make_shared<MiniMaxH3VAE::MiniMaxH3VideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
weight_manager);
} else if (sd_version_is_mage_flow(vae_version)) {
return std::make_shared<MageVAE::MageVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
weight_manager);
} else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
false,
vae_version,
weight_manager);
} else if (sd_version_uses_wan_vae(vae_version)) {
return std::make_shared<WAN::WanVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
false,
vae_version,
weight_manager);
} else {
auto model = std::make_shared<AutoEncoderKL>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
false,
false,
vae_version,
weight_manager);
if (sd_version_is_sdxl(version) &&
(strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) {
float vae_conv_2d_scale = 1.f / 32.f;
LOG_WARN(
"No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, "
"using Conv2D scale %.3f",
vae_conv_2d_scale);
model->set_conv2d_scale(vae_conv_2d_scale);
}
return model;
}
};
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
LOG_INFO("using FakeVAE");
result.vae = std::make_shared<FakeVAE>(version,
ctx.backends.runtime_backend(SDBackendModule::VAE),
weight_manager);
} else if (options.use_tae && !options.tae_preview_only) {
LOG_INFO("using TAE for encoding / decoding");
result.vae = create_tae(false);
} else {
LOG_INFO("using VAE for encoding / decoding");
result.vae = create_vae();
if (options.use_tae && options.tae_preview_only) {
LOG_INFO("using TAE for preview");
result.preview = create_tae(true);
}
}
if (options.use_audio_vae) {
if (sd_version_is_minimax_h3(version)) {
result.audio = std::make_shared<MiniMaxH3::AudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"",
weight_manager);
} else {
result.audio = std::make_shared<LTXV::LTXAudioVAERunner>(ctx.backends.runtime_backend(SDBackendModule::VAE),
tensor_storage_map,
"",
weight_manager);
}
}
if (sd_ctx_params->vae_conv_direct) {
LOG_INFO("Using Conv2d direct in the vae model");
result.vae->set_conv2d_direct_enabled(true);
if (result.preview) {
result.preview->set_conv2d_direct_enabled(true);
}
}
runners = std::move(result);
return true;
}
bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner) {
const auto* sd_ctx_params = &ctx.params;
const auto& tensor_storage_map = ctx.tensor_storage_map;
const auto version = ctx.version;
const auto& weight_manager = ctx.weight_manager;
if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) {
return false;
}
auto control_net = std::make_shared<ControlNet>(ctx.backends.runtime_backend(SDBackendModule::CONTROL_NET),
tensor_storage_map,
version,
"",
weight_manager);
if (sd_ctx_params->diffusion_conv_direct) {
LOG_INFO("Using Conv2d direct in the control net");
control_net->set_conv2d_direct_enabled(true);
}
runner = std::move(control_net);
return true;
}
bool build_extension_runners(const GenerationExtensionInitContext& ctx,
std::vector<std::shared_ptr<GenerationExtension>>& extensions) {
std::vector<std::shared_ptr<GenerationExtension>> result;
for (auto extension : {create_photomaker_extension(), create_pulid_extension()}) {
if (!extension->init(ctx)) {
return false;
}
if (extension->is_enabled()) {
result.push_back(std::move(extension));
}
}
extensions = std::move(result);
return true;
}
} // namespace sd::model_builders

63
src/model_builders.h Normal file
View File

@ -0,0 +1,63 @@
#ifndef __SD_MODEL_BUILDERS_H__
#define __SD_MODEL_BUILDERS_H__
#include <memory>
#include <vector>
#include "model.h"
#include "stable-diffusion.h"
class SDBackendManager;
struct DeviceResidencyManager;
struct Conditioner;
struct FrozenCLIPVisionEmbedder;
struct DiffusionModelRunner;
struct VAE;
struct AudioVAERunner;
struct ControlNet;
struct GenerationExtension;
struct GenerationExtensionInitContext;
namespace IPAdapter {
struct IPAdapterRunner;
}
namespace sd::model_builders {
struct Context {
const sd_ctx_params_t& params;
SDVersion version;
const String2TensorStorage& tensor_storage_map;
SDBackendManager& backends;
std::shared_ptr<DeviceResidencyManager> weight_manager;
};
struct CoreRunners {
std::shared_ptr<Conditioner> conditioner;
std::shared_ptr<DiffusionModelRunner> diffusion;
std::shared_ptr<DiffusionModelRunner> high_noise_diffusion;
std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;
std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
};
struct VAEOptions {
bool use_tae = false;
bool tae_preview_only = false;
bool use_audio_vae = false;
bool external_vae_is_invalid = false;
};
struct VAERunners {
std::shared_ptr<VAE> vae;
std::shared_ptr<VAE> preview;
std::shared_ptr<AudioVAERunner> audio;
};
bool build_core_runners(const Context& ctx, CoreRunners& runners);
bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners);
bool build_control_net_runner(const Context& ctx, std::shared_ptr<ControlNet>& runner);
bool build_extension_runners(const GenerationExtensionInitContext& ctx,
std::vector<std::shared_ptr<GenerationExtension>>& extensions);
} // namespace sd::model_builders
#endif // __SD_MODEL_BUILDERS_H__

61
src/model_component.h Normal file
View File

@ -0,0 +1,61 @@
#ifndef __SD_MODEL_COMPONENT_H__
#define __SD_MODEL_COMPONENT_H__
enum class ModelComponent {
Conditioner,
Diffusion,
HighNoiseDiffusion,
CLIPVision,
IPAdapter,
VAE,
PreviewVAE,
AudioVAE,
ControlNet,
PhotoMaker,
PuLID,
LoRA,
Upscaler,
Detector,
LatentUpsampler,
Count,
};
inline const char* model_component_name(ModelComponent component) {
switch (component) {
case ModelComponent::Conditioner:
return "Conditioner model";
case ModelComponent::Diffusion:
return "Diffusion model";
case ModelComponent::HighNoiseDiffusion:
return "High noise diffusion model";
case ModelComponent::CLIPVision:
return "CLIP vision";
case ModelComponent::IPAdapter:
return "IP-Adapter";
case ModelComponent::VAE:
return "VAE";
case ModelComponent::PreviewVAE:
return "preview VAE";
case ModelComponent::AudioVAE:
return "audio VAE";
case ModelComponent::ControlNet:
return "ControlNet";
case ModelComponent::PhotoMaker:
return "photomaker";
case ModelComponent::PuLID:
return "pulid";
case ModelComponent::LoRA:
return "LoRA";
case ModelComponent::Upscaler:
return "ESRGAN";
case ModelComponent::Detector:
return "YOLOv8";
case ModelComponent::LatentUpsampler:
return "LTX latent upsampler";
case ModelComponent::Count:
break;
}
return "unknown";
}
#endif // __SD_MODEL_COMPONENT_H__

39
src/model_io/safetensors_io.cpp
View File

@ -144,7 +144,8 @@ static bool read_comfy_quant_config(std::ifstream& file,
bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
std::string* error,
std::map<std::string, std::string>* metadata) {
std::map<std::string, std::string>* metadata,
std::set<std::string>* tensor_names) {
std::ifstream file(file_path, std::ios::binary);
if (!file.is_open()) {
set_error(error, "failed to open '" + file_path + "'");
@ -246,10 +247,6 @@ bool read_safetensors_file(const std::string& file_path,
std::string dtype = tensor_info["dtype"];
nlohmann::json shape = tensor_info["shape"];
if (dtype == "U8") {
continue;
}
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
size_t end = tensor_info["data_offsets"][1].get<size_t>();
if (begin > end || end > file_size_ - data_start) {
@ -257,6 +254,26 @@ bool read_safetensors_file(const std::string& file_path,
return false;
}
if (tensor_names != nullptr) {
tensor_names->insert(name);
}
if (dtype == "U8") {
uint64_t bytes = 1;
for (const auto& dimension : shape) {
const int64_t size = dimension.get<int64_t>();
if (size < 0 || (bytes != 0 && static_cast<uint64_t>(size) > UINT64_MAX / bytes)) {
set_error(error, "invalid dimensions for tensor '" + name + "'");
return false;
}
bytes *= size;
}
if (bytes != end - begin) {
set_error(error, "size mismatch for tensor '" + name + "'");
return false;
}
continue;
}
ggml_type type = safetensors_dtype_to_ggml_type(dtype);
if (type == GGML_TYPE_COUNT) {
set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
@ -270,8 +287,20 @@ bool read_safetensors_file(const std::string& file_path,
int n_dims = (int)shape.size();
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
uint64_t elements = 1;
for (int i = 0; i < n_dims; i++) {
ne[i] = shape[i].get<int64_t>();
if (ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(ne[i]) > INT64_MAX / elements)) {
set_error(error, "invalid dimensions for tensor '" + name + "'");
return false;
}
elements *= ne[i];
}
const uint64_t storage_size = ggml_type_size(type) * ((dtype == "F64" || dtype == "I64") ? 2 : 1);
if (elements % ggml_blck_size(type) != 0 ||
elements / ggml_blck_size(type) > INT64_MAX / storage_size) {
set_error(error, "invalid storage size for tensor '" + name + "'");
return false;
}
if (n_dims == 5) {

4
src/model_io/safetensors_io.h
View File

@ -2,6 +2,7 @@
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
#include <map>
#include <set>
#include <string>
#include <vector>
@ -12,7 +13,8 @@ bool is_safetensors_file(const std::string& file_path);
bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
std::string* error = nullptr,
std::map<std::string, std::string>* metadata = nullptr);
std::map<std::string, std::string>* metadata = nullptr,
std::set<std::string>* tensor_names = nullptr);
bool read_safetensors_index_file(const std::string& file_path,
std::vector<std::string>& shard_paths,
std::string* error = nullptr);

8
src/model_io/tensor_storage.h
View File

@ -28,9 +28,11 @@ struct TensorStorage {
int n_dims = 0;
std::string storage_key;
size_t file_index = 0;
int index_in_zip = -1; // >= means stored in a zip file
uint64_t offset = 0; // offset in file
size_t file_index = 0;
uint64_t file_id = 0;
uint64_t file_revision = 0;
int index_in_zip = -1; // >= means stored in a zip file
uint64_t offset = 0; // offset in file
TensorStorage() = default;

113
src/model_loader.cpp
View File

@ -9,6 +9,7 @@
#include <mutex>
#include <regex>
#include <set>
#include <stdexcept>
#include <string>
#include <thread>
#include <unordered_map>
@ -27,6 +28,7 @@
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "json.hpp"
#include "zip.h"
#include "name_conversion.h"
@ -151,15 +153,19 @@ ModelLoader::ModelLoader()
}
size_t ModelLoader::add_file_path(const std::string& file_path) {
if (model_files_processed) {
file_data.clear();
model_files_processed = false;
auto it = std::find(file_paths_.begin(), file_paths_.end(), file_path);
if (it != file_paths_.end()) {
return static_cast<size_t>(it - file_paths_.begin());
}
invalidate_file_data();
file_paths_.push_back(file_path);
return file_paths_.size() - 1;
}
void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
if (tensor_storage_map.count(tensor_storage.name) != 0) {
throw std::runtime_error("duplicate tensor in model source: " + tensor_storage.name);
}
tensor_storage_map[tensor_storage.name] = tensor_storage;
}
@ -169,6 +175,15 @@ void ModelLoader::set_n_threads(int n_threads) {
}
bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
return add_file(file_path, prefix);
}
bool ModelLoader::parse_file(const std::string& file_path, const std::string& prefix) {
FileStamp stamp;
if (!read_file_stamp(file_path, stamp)) {
return false;
}
parsed_dependencies_.push_back(stamp);
if (is_directory(file_path)) {
LOG_INFO("load %s using diffusers format", file_path.c_str());
return init_from_diffusers_file(file_path, prefix);
@ -198,17 +213,11 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
}
void ModelLoader::convert_tensors_name() {
SDVersion version = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
String2TensorStorage new_map;
for (auto& [_, tensor_storage] : tensor_storage_map) {
auto new_name = convert_tensor_name(tensor_storage.name, version);
// LOG_VERBOSE("%s -> %s", tensor_storage.name.c_str(), new_name.c_str());
tensor_storage.name = new_name;
new_map[new_name] = std::move(tensor_storage);
if (names_converted_) {
return;
}
tensor_storage_map.swap(new_map);
names_converted_ = true;
rebuild_catalog();
}
bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, const std::string& prefix, SDVersion version) {
@ -257,7 +266,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
std::vector<TensorStorage> tensor_storages;
std::string error;
if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_, &parsed_tensor_names_[file_path])) {
LOG_ERROR("%s", error.c_str());
return false;
}
@ -293,7 +302,26 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
}
for (const std::string& shard_path : shard_paths) {
if (!init_from_file(shard_path, prefix)) {
if (!parse_file(shard_path, prefix)) {
return false;
}
}
std::ifstream index_file(file_path);
const auto index = nlohmann::json::parse(index_file);
for (const auto& entry : index.at("weight_map").items()) {
const auto expected = (std::filesystem::u8path(file_path).parent_path() /
std::filesystem::u8path(entry.value().get<std::string>()))
.lexically_normal();
bool found = false;
for (const auto& shard : parsed_tensor_names_) {
if (std::filesystem::u8path(shard.first).lexically_normal() == expected) {
found = shard.second.count(entry.key()) != 0;
break;
}
}
if (!found) {
LOG_ERROR("safetensors index tensor '%s' is missing from its declared shard", entry.key().c_str());
return false;
}
}
@ -369,25 +397,23 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
std::string clip_path = path_join(file_path, "text_encoder/model.safetensors");
std::string clip_g_path = path_join(file_path, "text_encoder_2/model.safetensors");
if (!init_from_safetensors_file(unet_path, "unet.")) {
if (!parse_file(unet_path, prefix + "unet.")) {
return false;
}
if (!init_from_safetensors_file(vae_path, "vae.")) {
LOG_WARN("Couldn't find working VAE in %s", file_path.c_str());
// return false;
if (file_exists(vae_path) && !parse_file(vae_path, prefix + "vae.")) {
return false;
}
if (!init_from_safetensors_file(clip_path, "te.")) {
LOG_WARN("Couldn't find working text encoder in %s", file_path.c_str());
// return false;
if (file_exists(clip_path) && !parse_file(clip_path, prefix + "te.")) {
return false;
}
if (!init_from_safetensors_file(clip_g_path, "te.1.")) {
LOG_VERBOSE("Couldn't find working second text encoder in %s", file_path.c_str());
if (file_exists(clip_g_path) && !parse_file(clip_g_path, prefix + "te.1.")) {
return false;
}
return true;
}
SDVersion ModelLoader::get_sd_version() {
SDVersion ModelLoader::get_sd_version() const {
TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
bool has_multiple_encoders = false;
@ -623,7 +649,7 @@ SDVersion ModelLoader::get_sd_version() {
return VERSION_COUNT;
}
std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) {
@ -640,7 +666,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
return wtype_stat;
}
std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) {
@ -664,7 +690,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
return wtype_stat;
}
std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) {
@ -685,7 +711,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
return wtype_stat;
}
std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() {
std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) {
@ -743,9 +769,12 @@ TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules) {
}
void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_rules) {
auto map_rules = parse_tensor_type_rules(tensor_type_rules);
wtype_override_ = wtype;
tensor_type_rules_ = tensor_type_rules;
auto map_rules = parse_tensor_type_rules(tensor_type_rules);
for (auto& [name, tensor_storage] : tensor_storage_map) {
ggml_type dst_type = wtype;
tensor_storage.expected_type = GGML_TYPE_COUNT;
ggml_type dst_type = wtype;
for (const auto& tensor_type_rule : map_rules) {
std::regex pattern(tensor_type_rule.first);
if (std::regex_search(name, pattern)) {
@ -761,6 +790,8 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
}
tensor_storage.expected_type = dst_type;
}
invalidate_file_data();
++revision_;
}
void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
@ -829,6 +860,13 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors,
bool writable_mmap) {
std::set<std::string> names;
for (const auto& entry : tensors) {
names.insert(entry.first);
}
if (!validate_sources(&names)) {
return {};
}
process_model_files(true, writable_mmap);
std::vector<MmapTensorStore> result;
@ -919,6 +957,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool enable_mmap,
const std::set<std::string>* target_tensor_names,
bool log_progress) {
if (!validate_sources(target_tensor_names)) {
return false;
}
process_model_files(enable_mmap, false);
std::atomic<int64_t> read_time_ms(0);
@ -1242,7 +1283,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
(convert_time_ms.load() / (float)last_n_threads) / 1000.f,
(copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
}
return success;
return success && validate_sources(target_tensor_names);
}
bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor) {
@ -1259,7 +1300,9 @@ bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor*
return true;
}
if (current_tensor_storage.file_index != tensor_storage.file_index ||
if (current_tensor_storage.file_id != tensor_storage.file_id ||
current_tensor_storage.file_revision != tensor_storage.file_revision ||
current_tensor_storage.file_index != tensor_storage.file_index ||
current_tensor_storage.offset != tensor_storage.offset ||
current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) {
LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
@ -1440,7 +1483,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
return true;
}
bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const {
const std::string& name = tensor_storage.name;
if (tensor_storage.is_int8_tensorwise) {
return false;
@ -1478,7 +1521,7 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
return false;
}
int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) {
int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) const {
size_t alignment = 128;
if (backend != nullptr) {
alignment = ggml_backend_get_alignment(backend);

66
src/model_loader.h
View File

@ -2,6 +2,7 @@
#define __MODEL_LOADER_H__
#include <cstdint>
#include <filesystem>
#include <map>
#include <memory>
#include <set>
@ -30,6 +31,46 @@ struct MmapTensorStore {
bool is_unused_tensor(const std::string& name);
class ModelLoader {
public:
using FileId = uint64_t;
using FileVersions = std::map<FileId, uint64_t>;
enum class FileScope { Catalog,
Isolated };
private:
struct FileStamp {
std::string path;
uintmax_t size = 0;
std::filesystem::file_time_type modified;
};
struct FileRecord {
FileId id = 0;
uint64_t revision = 0;
std::string path;
std::string prefix;
FileScope scope = FileScope::Catalog;
std::vector<FileStamp> dependencies;
String2TensorStorage tensors;
std::map<std::string, std::string> metadata;
};
std::vector<FileRecord> files_;
uint64_t revision_ = 0;
bool names_converted_ = false;
ggml_type wtype_override_ = GGML_TYPE_COUNT;
std::string tensor_type_rules_;
std::vector<FileStamp> parsed_dependencies_;
std::map<std::string, std::set<std::string>> parsed_tensor_names_;
static bool read_file_stamp(const std::string& path, FileStamp& stamp);
static bool file_unchanged(const FileStamp& stamp);
bool parse_file(const std::string& path, const std::string& prefix);
bool add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope);
ModelLoader file_reader(FileId id, SDVersion version) const;
void rebuild_catalog();
void invalidate_file_data();
protected:
SDVersion version_ = VERSION_COUNT;
std::vector<std::string> file_paths_;
@ -52,16 +93,27 @@ protected:
public:
ModelLoader();
bool add_file(const std::string& path, const std::string& prefix = "", FileId* id = nullptr, bool force = false, FileScope scope = FileScope::Catalog);
bool del_file(FileId id);
uint64_t file_revision(FileId id) const;
std::string file_path(FileId id) const;
String2TensorStorage file_tensors(FileId id, SDVersion version) const;
bool load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap = false) const;
bool refresh_files(bool include_isolated = true);
bool files_changed(bool& changed, bool include_isolated = true) const;
bool validate_sources(const std::set<std::string>* tensor_names = nullptr) const;
uint64_t revision() const { return revision_; }
FileVersions file_versions(const std::vector<std::string>& prefixes = {}) const;
bool init_from_file(const std::string& file_path, const std::string& prefix = "");
void convert_tensors_name();
bool init_from_file_and_convert_name(const std::string& file_path,
const std::string& prefix = "",
SDVersion version = VERSION_COUNT);
SDVersion get_sd_version();
std::map<ggml_type, uint32_t> get_wtype_stat();
std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
std::map<ggml_type, uint32_t> get_vae_wtype_stat();
SDVersion get_sd_version() const;
std::map<ggml_type, uint32_t> get_wtype_stat() const;
std::map<ggml_type, uint32_t> get_conditioner_wtype_stat() const;
std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat() const;
std::map<ggml_type, uint32_t> get_vae_wtype_stat() const;
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
@ -92,8 +144,8 @@ public:
return names;
}
bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const;
int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT) const;
~ModelLoader() = default;
};

336
src/model_loader_files.cpp Normal file
View File

@ -0,0 +1,336 @@
#include "model_loader.h"
#include <algorithm>
#include <atomic>
#include <limits>
#include "core/util.h"
#include "name_conversion.h"
static uint64_t next_source_revision() {
static std::atomic<uint64_t> revision{0};
return revision.fetch_add(1, std::memory_order_relaxed) + 1;
}
bool ModelLoader::read_file_stamp(const std::string& path, FileStamp& stamp) {
std::error_code error;
const auto file_path = std::filesystem::u8path(path);
stamp.path = path;
stamp.size = 0;
stamp.modified = std::filesystem::last_write_time(file_path, error);
if (!error && std::filesystem::is_regular_file(file_path, error)) {
stamp.size = std::filesystem::file_size(file_path, error);
}
if (error) {
LOG_ERROR("cannot inspect model source '%s': %s", path.c_str(), error.message().c_str());
return false;
}
return true;
}
bool ModelLoader::file_unchanged(const FileStamp& stamp) {
std::error_code error;
if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error)) {
return false;
}
FileStamp current;
return read_file_stamp(stamp.path, current) &&
current.size == stamp.size && current.modified == stamp.modified;
}
void ModelLoader::invalidate_file_data() {
file_data.clear();
model_files_processed = false;
}
void ModelLoader::rebuild_catalog() {
tensor_storage_map.clear();
metadata_.clear();
for (const auto& file : files_) {
if (file.scope == FileScope::Isolated)
continue;
for (const auto& entry : file.tensors) {
tensor_storage_map[entry.first] = entry.second;
}
for (const auto& entry : file.metadata) {
metadata_[entry.first] = entry.second;
}
}
if (names_converted_) {
const SDVersion version = version_ == VERSION_COUNT ? get_sd_version() : version_;
tensor_storage_map.clear();
for (const auto& file : files_) {
if (file.scope == FileScope::Isolated)
continue;
for (const auto& entry : file.tensors) {
TensorStorage tensor = entry.second;
tensor.name = convert_tensor_name(tensor.name, version);
tensor_storage_map[tensor.name] = std::move(tensor);
}
}
}
std::set<size_t> used_files;
for (const auto& file : files_) {
for (const auto& entry : file.tensors) {
used_files.insert(entry.second.file_index);
}
}
for (size_t i = 0; i < file_paths_.size(); ++i) {
if (used_files.count(i) == 0) {
file_paths_[i].clear();
}
}
set_wtype_override(wtype_override_, tensor_type_rules_);
}
bool ModelLoader::add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
FileStamp root;
if (!read_file_stamp(path, root)) {
return false;
}
auto existing = std::find_if(files_.begin(), files_.end(), [&](const FileRecord& file) {
return file.path == root.path && file.prefix == prefix && file.scope == scope;
});
if (existing != files_.end() && !force &&
std::all_of(existing->dependencies.begin(), existing->dependencies.end(), file_unchanged)) {
if (id != nullptr) {
*id = existing->id;
}
return true;
}
ModelLoader parsed;
try {
if (!parsed.parse_file(root.path, prefix)) {
return false;
}
} catch (const std::exception& error) {
LOG_ERROR("invalid model source '%s': %s", path.c_str(), error.what());
return false;
}
std::vector<size_t> file_indices;
std::vector<FileStamp> physical_files;
for (const auto& physical_path : parsed.file_paths_) {
FileStamp stamp;
if (!read_file_stamp(physical_path, stamp)) {
return false;
}
parsed.parsed_dependencies_.push_back(stamp);
file_indices.push_back(add_file_path(stamp.path));
physical_files.push_back(std::move(stamp));
}
for (auto& entry : parsed.tensor_storage_map) {
auto& tensor = entry.second;
// Pickle preserves rank-zero scalars; GGML uses a one-element dimension.
if (tensor.n_dims == 0) {
tensor.n_dims = 1;
}
if (tensor.n_dims < 1 || tensor.n_dims > SD_MAX_DIMS || tensor.type < 0 ||
tensor.type >= GGML_TYPE_COUNT || tensor.file_index >= parsed.file_paths_.size()) {
LOG_ERROR("invalid tensor metadata for '%s'", tensor.name.c_str());
return false;
}
uint64_t elements = 1;
for (int i = 0; i < tensor.n_dims; ++i) {
if (tensor.ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(tensor.ne[i]) > INT64_MAX / elements)) {
LOG_ERROR("invalid tensor dimensions for '%s'", tensor.name.c_str());
return false;
}
elements *= tensor.ne[i];
}
const uint64_t block_size = ggml_blck_size(tensor.type);
const uint64_t type_size = ggml_type_size(tensor.type) * ((tensor.is_f64 || tensor.is_i64) ? 2 : 1);
if (block_size == 0 || type_size == 0 || elements % block_size != 0 || elements / block_size > INT64_MAX / type_size) {
LOG_ERROR("invalid tensor storage size for '%s'", tensor.name.c_str());
return false;
}
if (tensor.index_in_zip < 0) {
const auto& stamp = physical_files[tensor.file_index];
if (tensor.offset > stamp.size || elements / block_size * type_size > stamp.size - tensor.offset) {
LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str());
return false;
}
}
}
if (!std::all_of(parsed.parsed_dependencies_.begin(), parsed.parsed_dependencies_.end(), file_unchanged)) {
LOG_ERROR("model source changed while reading metadata: '%s'", path.c_str());
return false;
}
FileRecord record;
// Snapshots and independently created loaders must never alias different versions.
record.revision = next_source_revision();
record.id = existing == files_.end() ? record.revision : existing->id;
++revision_;
record.path = root.path;
record.prefix = prefix;
record.scope = scope;
std::set<std::string> seen_dependencies;
for (auto& stamp : parsed.parsed_dependencies_) {
if (seen_dependencies.insert(stamp.path).second) {
record.dependencies.push_back(std::move(stamp));
}
}
record.metadata = std::move(parsed.metadata_);
record.tensors = std::move(parsed.tensor_storage_map);
for (auto& entry : record.tensors) {
entry.second.file_index = file_indices[entry.second.file_index];
entry.second.file_id = record.id;
entry.second.file_revision = record.revision;
}
if (id != nullptr) {
*id = record.id;
}
if (existing == files_.end()) {
files_.push_back(std::move(record));
} else {
*existing = std::move(record);
}
rebuild_catalog();
return true;
}
bool ModelLoader::add_file(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope) {
ModelLoader candidate = *this;
FileId added_id = 0;
if (!candidate.add_file_impl(path, prefix, &added_id, force, scope)) {
return false;
}
*this = std::move(candidate);
if (id != nullptr) {
*id = added_id;
}
return true;
}
bool ModelLoader::del_file(FileId id) {
auto it = std::find_if(files_.begin(), files_.end(), [id](const FileRecord& file) { return file.id == id; });
if (it == files_.end()) {
return false;
}
files_.erase(it);
++revision_;
rebuild_catalog();
return true;
}
bool ModelLoader::files_changed(bool& changed, bool include_isolated) const {
changed = false;
for (const auto& file : files_) {
if (!include_isolated && file.scope == FileScope::Isolated)
continue;
for (const auto& stamp : file.dependencies) {
std::error_code error;
if (!std::filesystem::exists(std::filesystem::u8path(stamp.path), error) && !error) {
// An updated index may no longer reference this dependency.
changed = true;
continue;
}
FileStamp current;
if (!read_file_stamp(stamp.path, current)) {
return false;
}
changed |= current.size != stamp.size || current.modified != stamp.modified;
}
}
return true;
}
bool ModelLoader::refresh_files(bool include_isolated) {
bool changed;
if (!files_changed(changed, include_isolated)) {
return false;
}
if (!changed) {
return true;
}
ModelLoader candidate = *this;
for (const auto& file : files_) {
if (!include_isolated && file.scope == FileScope::Isolated)
continue;
if (!candidate.add_file_impl(file.path, file.prefix, nullptr, false, file.scope)) {
return false;
}
}
*this = std::move(candidate);
return true;
}
bool ModelLoader::validate_sources(const std::set<std::string>* tensor_names) const {
std::set<FileId> required;
if (tensor_names != nullptr) {
for (const auto& name : *tensor_names) {
auto it = tensor_storage_map.find(name);
if (it != tensor_storage_map.end()) {
required.insert(it->second.file_id);
}
}
}
for (const auto& file : files_) {
if (tensor_names != nullptr && required.count(file.id) == 0) {
continue;
}
if (!std::all_of(file.dependencies.begin(), file.dependencies.end(), file_unchanged)) {
LOG_ERROR("model source changed; refresh it before execution: '%s'", file.path.c_str());
return false;
}
}
return true;
}
ModelLoader::FileVersions ModelLoader::file_versions(const std::vector<std::string>& prefixes) const {
FileVersions versions;
for (const auto& entry : tensor_storage_map) {
if (prefixes.empty() || std::any_of(prefixes.begin(), prefixes.end(), [&](const std::string& prefix) {
return starts_with(entry.first, prefix);
})) {
versions[entry.second.file_id] = entry.second.file_revision;
}
}
return versions;
}
uint64_t ModelLoader::file_revision(FileId id) const {
for (const auto& file : files_) {
if (file.id == id)
return file.revision;
}
return 0;
}
std::string ModelLoader::file_path(FileId id) const {
for (const auto& file : files_) {
if (file.id == id)
return file.path;
}
return {};
}
ModelLoader ModelLoader::file_reader(FileId id, SDVersion version) const {
ModelLoader reader;
reader.file_paths_ = file_paths_;
reader.n_threads_ = n_threads_;
reader.version_ = version;
reader.names_converted_ = true;
for (const auto& file : files_) {
if (file.id == id) {
reader.files_.push_back(file);
reader.files_.back().scope = FileScope::Catalog;
break;
}
}
reader.rebuild_catalog();
return reader;
}
String2TensorStorage ModelLoader::file_tensors(FileId id, SDVersion version) const {
return file_reader(id, version).tensor_storage_map;
}
bool ModelLoader::load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap) const {
if (file_revision(id) == 0)
return false;
auto reader = file_reader(id, version);
return reader.load_tensors(callback, use_mmap, &names, false);
}

323
src/model_manager.cpp
View File

@ -4,6 +4,7 @@
#include <cstdint>
#include <iterator>
#include <mutex>
#include <tuple>
#include <unordered_set>
#include "core/ggml_extend_backend.h"
@ -26,7 +27,8 @@ static bool lora_specs_equal(const std::vector<ModelManager::LoraSpec>& lhs,
lhs[i].multiplier != rhs[i].multiplier ||
lhs[i].is_high_noise != rhs[i].is_high_noise ||
lhs[i].tensor_name_prefix_filter != rhs[i].tensor_name_prefix_filter ||
lhs[i].required != rhs[i].required) {
lhs[i].required != rhs[i].required ||
lhs[i].file_id != rhs[i].file_id || lhs[i].file_revision != rhs[i].file_revision) {
return false;
}
}
@ -104,25 +106,61 @@ void ModelManager::set_common_ignore_tensors(std::set<std::string> ignore_tensor
common_ignore_tensors_ = std::move(ignore_tensors);
}
void ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
if (loras.empty() && loras_.empty()) {
lora_version_ = version;
return;
bool ModelManager::prepare_lora_sources(std::vector<LoraSpec>& loras) {
ModelLoader candidate = model_loader_;
std::vector<LoraSpec> resolved;
std::set<ModelLoader::FileId> sources;
for (auto spec : loras) {
const std::string prefix = spec.is_high_noise ? "lora.model.high_noise_" : "lora.";
if (!candidate.add_file(spec.path, prefix, &spec.file_id, false, ModelLoader::FileScope::Isolated)) {
if (spec.required)
return false;
LOG_WARN("cannot register LoRA source '%s'", spec.path.c_str());
continue;
}
spec.file_revision = candidate.file_revision(spec.file_id);
sources.insert(spec.file_id);
resolved.push_back(std::move(spec));
}
if (lora_version_ == version && lora_specs_equal(loras_, loras)) {
return;
for (auto id : lora_sources_) {
if (sources.count(id) == 0)
candidate.del_file(id);
}
if (!set_loader(std::move(candidate)))
return false;
lora_sources_ = std::move(sources);
loras = std::move(resolved);
return true;
}
bool ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
if (std::any_of(loras.begin(), loras.end(), [](const LoraSpec& spec) { return spec.file_id == 0; }) &&
!prepare_lora_sources(loras))
return false;
for (auto& spec : loras) {
spec.file_revision = model_loader_.file_revision(spec.file_id);
if (spec.file_revision == 0)
return false;
}
if (lora_version_ == version && lora_specs_equal(loras_, loras))
return true;
if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
return state->pin_count != 0;
})) {
LOG_ERROR("cannot change LoRA configuration during execution");
return false;
}
loras_ = std::move(loras);
lora_version_ = version;
current_lora_epoch_++;
reset_lora_applied_params();
return true;
}
std::set<std::string> ModelManager::tensor_names() const {
std::set<std::string> names;
for (const auto& state : tensor_states_) {
if (state != nullptr) {
if (state != nullptr && state->component != ModelComponent::LoRA) {
names.insert(state->name);
}
}
@ -171,7 +209,7 @@ ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState
return state.split_buffer_type;
}
bool ModelManager::register_param_tensors(const std::string& desc,
bool ModelManager::register_param_tensors(ModelComponent component,
std::map<std::string, ggml_tensor*> tensors,
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
@ -179,15 +217,20 @@ bool ModelManager::register_param_tensors(const std::string& desc,
size_t* registered_tensor_size,
bool allow_split_buffer,
bool params_follow_compute_backend,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops,
ModelLoader::FileId source_file,
SDVersion source_version) {
if (component == ModelComponent::Count) {
LOG_ERROR("model manager tensor component is invalid");
return false;
}
if (registered_tensor_size != nullptr) {
*registered_tensor_size += estimate_tensors_size(tensors);
}
const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{};
const auto& sources = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map();
std::unordered_set<ggml_tensor*> new_tensors;
std::vector<std::unique_ptr<TensorState>> new_states;
new_states.reserve(tensors.size());
@ -197,16 +240,23 @@ bool ModelManager::register_param_tensors(const std::string& desc,
if (tensor == nullptr) {
continue;
}
if (tensor_states_by_name_.find(name) != tensor_states_by_name_.end()) {
if (tensor_states_by_tensor_.count(tensor) != 0 || !new_tensors.insert(tensor).second) {
LOG_ERROR("model manager tensor name '%s' is already registered", name.c_str());
return false;
}
ggml_set_name(tensor, name.c_str());
auto state = std::make_unique<TensorState>();
state->name = name;
state->tensor = tensor;
state->desc = desc;
auto state = std::make_unique<TensorState>();
state->name = name;
state->tensor = tensor;
state->component = component;
state->source_file = source_file;
state->source_version = source_version;
auto source = sources.find(name);
if (source != sources.end()) {
state->source = source->second;
state->has_source = true;
}
state->residency_mode = residency_mode;
state->compute_backend = compute_backend;
state->params_backend = params_backend;
@ -225,31 +275,45 @@ bool ModelManager::register_param_tensors(const std::string& desc,
}
for (auto& state : new_states) {
TensorState* registered_state = state.get();
tensor_states_by_name_[registered_state->name] = registered_state;
TensorState* registered_state = state.get();
tensor_states_by_tensor_[registered_state->tensor] = registered_state;
tensor_states_.push_back(std::move(state));
}
return true;
}
bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
if (desc.empty()) {
return true;
bool ModelManager::unregister_param_tensors(ModelComponent component, size_t* registered_tensor_size) {
std::unordered_set<TensorState*> states;
for (auto& state : tensor_states_) {
if (state->component == component)
states.insert(state.get());
}
return unregister_tensor_states(states, registered_tensor_size);
}
std::unordered_set<TensorState*> target_states;
bool ModelManager::unregister_param_tensors(const std::vector<ggml_tensor*>& tensors) {
std::unordered_set<TensorState*> states;
for (auto tensor : tensors) {
auto found = tensor_states_by_tensor_.find(tensor);
if (found != tensor_states_by_tensor_.end())
states.insert(found->second);
}
return unregister_tensor_states(states, nullptr);
}
bool ModelManager::unregister_tensor_states(const std::unordered_set<TensorState*>& target_states,
size_t* registered_tensor_size) {
size_t released_size = 0;
for (auto& state : tensor_states_) {
if (state == nullptr || state->desc != desc) {
if (state == nullptr || target_states.count(state.get()) == 0) {
continue;
}
if (state->pin_count > 0) {
LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
desc.c_str(),
model_component_name(state->component),
state->name.c_str());
return false;
}
target_states.insert(state.get());
if (state->tensor != nullptr) {
released_size += ggml_nbytes(state->tensor);
}
@ -260,7 +324,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
}
clear_all_prefetched_params();
release_compute_staging_blocks(false);
release_compute_staging_blocks(false, &target_states);
std::vector<ParamsStorageBlock*> storage_blocks_to_release;
std::unordered_set<TensorState*> affected_storage_states;
@ -292,7 +356,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
}
if (state->pin_count > 0 || state->staged_to_compute_backend) {
LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
desc.c_str(),
model_component_name(state->component),
state->name.c_str());
return false;
}
@ -305,9 +369,9 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg
}
}
for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
for (auto it = tensor_states_by_tensor_.begin(); it != tensor_states_by_tensor_.end();) {
if (target_states.count(it->second) > 0) {
it = tensor_states_by_name_.erase(it);
it = tensor_states_by_tensor_.erase(it);
} else {
++it;
}
@ -559,19 +623,24 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
}
bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states) {
if (loras_.empty()) {
if (loras_.empty() || applying_loras_)
return true;
}
applying_loras_ = true;
struct ApplyGuard {
bool& active;
~ApplyGuard() { active = false; }
} guard{applying_loras_};
struct LoraApplyGroup {
std::map<std::string, ggml_tensor*> model_tensors;
std::vector<TensorState*> states;
};
std::map<ggml_backend_t, LoraApplyGroup> groups;
using ApplyTarget = std::tuple<ggml_backend_t, ggml_backend_t, ResidencyMode>;
std::map<ApplyTarget, LoraApplyGroup> groups;
for (TensorState* state : states) {
if (state == nullptr || state->tensor == nullptr ||
should_ignore(*state) || is_optional_missing_tensor(state->name)) {
if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA ||
state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
continue;
}
if (state->applied_lora_epoch == current_lora_epoch_) {
@ -596,7 +665,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
return false;
}
LoraApplyGroup& group = groups[state->compute_backend];
LoraApplyGroup& group = groups[{state->compute_backend, state->params_backend, state->residency_mode}];
group.model_tensors[state->name] = state->tensor;
group.states.push_back(state);
}
@ -607,20 +676,20 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
std::set<std::string> all_tensor_names = tensor_names();
for (auto& group_pair : groups) {
ggml_backend_t compute_backend = group_pair.first;
ggml_backend_t compute_backend = std::get<0>(group_pair.first);
LoraApplyGroup& group = group_pair.second;
for (const LoraSpec& lora_spec : loras_) {
if (group.model_tensors.empty()) {
continue;
}
std::string id = lora_id(lora_spec);
auto lora = std::make_shared<LoraModel>(id,
compute_backend,
compute_backend,
lora_spec.path,
lora_spec.is_high_noise ? "model.high_noise_" : "",
lora_version_);
std::string id = lora_id(lora_spec);
const auto* target = group.states.front();
// The temporary runner is destroyed before this manager call returns.
auto borrowed_manager = std::shared_ptr<ModelManager>(this, [](ModelManager*) {});
auto lora = std::make_shared<LoraModel>(id, compute_backend, target->params_backend,
borrowed_manager, lora_spec.file_id, lora_version_,
target->residency_mode);
LoraModel::filter_t lora_tensor_filter = nullptr;
if (!lora_spec.tensor_name_prefix_filter.empty()) {
@ -628,7 +697,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
return starts_with(tensor_name, lora_spec.tensor_name_prefix_filter);
};
}
if (!lora->load_from_file(n_threads_, lora_tensor_filter)) {
if (!lora->init_params(n_threads_, lora_tensor_filter)) {
LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str());
if (lora_spec.required) {
return false;
@ -643,7 +712,8 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
continue;
}
lora->multiplier = lora_spec.multiplier;
lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false);
if (!lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false))
return false;
lora->release_loaded_tensors();
}
@ -657,12 +727,13 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
}
void ModelManager::reset_lora_applied_params() {
clear_all_prefetched_params();
release_compute_staging_blocks(true);
release_params_storage_blocks(true);
std::unordered_set<TensorState*> affected;
for (auto& state : tensor_states_) {
state->applied_lora_epoch = UINT64_MAX;
if (state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX) {
affected.insert(state.get());
}
}
invalidate_sources(affected);
}
bool ModelManager::should_ignore(const TensorState& state) const {
@ -684,21 +755,19 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
return true;
}
const auto& tensor_storage_map = model_loader_.get_tensor_storage_map();
auto ts_it = tensor_storage_map.find(state.name);
if (ts_it == tensor_storage_map.end()) {
LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str());
if (!state.has_source) {
LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
return false;
}
const TensorStorage& tensor_storage = ts_it->second;
const TensorStorage& tensor_storage = state.source;
if (state.tensor->ne[0] != tensor_storage.ne[0] ||
state.tensor->ne[1] != tensor_storage.ne[1] ||
state.tensor->ne[2] != tensor_storage.ne[2] ||
state.tensor->ne[3] != tensor_storage.ne[3]) {
LOG_ERROR(
"%s tensor '%s' has wrong shape in model metadata: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
state.desc.c_str(),
model_component_name(state.component),
state.name.c_str(),
(int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
(int)state.tensor->ne[0], (int)state.tensor->ne[1], (int)state.tensor->ne[2], (int)state.tensor->ne[3]);
@ -746,7 +815,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
}
bool ModelManager::can_mmap_storage(const TensorState& state) const {
if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
return false;
}
if (state.compute_backend == nullptr || state.params_backend == nullptr) {
@ -857,75 +926,55 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
}
bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
std::map<std::string, TensorState*> states_by_name;
std::set<std::string> target_tensor_names;
for (TensorState* state : states) {
if (state == nullptr) {
using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
std::map<ReadGroup, std::vector<ReadBatch>> groups;
for (auto* state : states) {
if (state == nullptr)
continue;
auto& batches = groups[{state->source_file, state->source_version}];
// The loader supplies one destination per name; only conflicting types need another batch.
auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
auto found = candidate.find(state->name);
return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type;
});
if (batch == batches.end()) {
batches.emplace_back();
batch = std::prev(batches.end());
}
states_by_name[state->name] = state;
target_tensor_names.insert(state->name);
(*batch)[state->name].push_back(state);
}
if (states_by_name.empty()) {
return true;
}
std::set<std::string> loaded_names;
std::mutex loaded_names_mutex;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
const std::string& name = tensor_storage.name;
*dst_tensor = nullptr;
auto state_it = states_by_name.find(name);
if (state_it == states_by_name.end()) {
return true;
for (auto& group : groups) {
for (auto& batch : group.second) {
std::set<std::string> names;
std::set<std::string> loaded;
std::mutex mutex;
for (const auto& entry : batch)
names.insert(entry.first);
auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
*dst = nullptr;
auto found = batch.find(source.name);
if (found == batch.end())
return true;
*dst = found->second.front()->tensor;
std::lock_guard<std::mutex> lock(mutex);
loaded.insert(source.name);
return true;
};
const auto file = group.first.first;
bool success = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names)
: model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_);
if (!success || loaded != names)
return false;
for (auto& entry : batch) {
auto* first = entry.second.front()->tensor;
for (auto* state : entry.second) {
if (state->tensor != first)
ggml_backend_tensor_copy(first, state->tensor);
state->loaded_to_params_backend = true;
}
}
}
TensorState* state = state_it->second;
if (state == nullptr || state->tensor == nullptr) {
LOG_ERROR("model manager tensor '%s' is null", name.c_str());
return false;
}
if (state->tensor->ne[0] != tensor_storage.ne[0] ||
state->tensor->ne[1] != tensor_storage.ne[1] ||
state->tensor->ne[2] != tensor_storage.ne[2] ||
state->tensor->ne[3] != tensor_storage.ne[3]) {
LOG_ERROR(
"model manager tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
name.c_str(),
(int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
(int)state->tensor->ne[0], (int)state->tensor->ne[1], (int)state->tensor->ne[2], (int)state->tensor->ne[3]);
return false;
}
{
std::lock_guard<std::mutex> lock(loaded_names_mutex);
loaded_names.insert(name);
}
*dst_tensor = state->tensor;
return true;
};
if (!model_loader_.load_tensors(on_new_tensor_cb, enable_mmap_, &target_tensor_names)) {
LOG_ERROR("model manager load tensors failed");
return false;
}
bool missing = false;
for (const auto& pair : states_by_name) {
const std::string& name = pair.first;
if (loaded_names.find(name) == loaded_names.end()) {
LOG_ERROR("model manager tensor '%s' was not loaded", name.c_str());
missing = true;
}
}
if (missing) {
return false;
}
for (const auto& pair : states_by_name) {
pair.second->loaded_to_params_backend = true;
}
return true;
}
@ -1138,6 +1187,14 @@ void ModelManager::release_all() {
release_params_storage_blocks(true);
}
ggml_tensor* ModelManager::resolve_param_tensor(ggml_tensor* tensor) const {
for (auto* current = tensor; current != nullptr; current = current->view_src) {
if (tensor_states_by_tensor_.count(current) != 0)
return current;
}
return nullptr;
}
bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
std::vector<TensorState*>& required_states,
ggml_backend_t compute_backend) const {
@ -1147,21 +1204,13 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
if (tensor == nullptr) {
continue;
}
const char* raw_name = ggml_get_name(tensor);
if (raw_name == nullptr || raw_name[0] == '\0') {
LOG_ERROR("model manager unnamed tensor is not registered");
return false;
}
auto state_it = tensor_states_by_name_.find(raw_name);
if (state_it == tensor_states_by_name_.end()) {
LOG_ERROR("model manager tensor '%s' is not registered", raw_name);
return false;
}
TensorState* state = state_it->second;
if (state == nullptr) {
LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name);
auto param = resolve_param_tensor(tensor);
auto found = tensor_states_by_tensor_.find(param);
if (found == tensor_states_by_tensor_.end()) {
LOG_ERROR("model manager tensor '%s' is not registered", ggml_get_name(tensor));
return false;
}
TensorState* state = found->second;
if ((compute_backend == nullptr || state->compute_backend == nullptr ||
state->compute_backend == compute_backend) &&
seen.insert(state).second) {
@ -1375,8 +1424,8 @@ bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
}
if (!apply_loras_to_params(required_states)) {
finish_compute_backend_usage(required_states);
release_compute_staging_blocks(false);
release_params_storage_blocks(false);
std::unordered_set<TensorState*> failed(required_states.begin(), required_states.end());
invalidate_sources(failed);
return false;
}
return true;

54
src/model_manager.h
View File

@ -10,6 +10,7 @@
#include <vector>
#include "device_residency_manager.h"
#include "model_component.h"
#include "model_loader.h"
class ModelManager : public DeviceResidencyManager {
@ -24,7 +25,9 @@ public:
float multiplier = 1.0f;
bool is_high_noise = false;
std::string tensor_name_prefix_filter;
bool required = false;
bool required = false;
ModelLoader::FileId file_id = 0;
uint64_t file_revision = 0;
};
private:
@ -32,8 +35,12 @@ private:
struct TensorState {
std::string name;
ggml_tensor* tensor = nullptr;
std::string desc;
ggml_tensor* tensor = nullptr;
ModelComponent component = ModelComponent::Count;
TensorStorage source;
bool has_source = false;
ModelLoader::FileId source_file = 0;
SDVersion source_version = VERSION_COUNT;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
@ -79,7 +86,7 @@ private:
ModelLoader model_loader_;
std::vector<std::unique_ptr<TensorState>> tensor_states_;
std::map<std::string, TensorState*> tensor_states_by_name_;
std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
@ -91,6 +98,8 @@ private:
bool warned_split_lora_skip_ = false;
std::set<std::string> common_ignore_tensors_;
std::vector<LoraSpec> loras_;
std::set<ModelLoader::FileId> lora_sources_;
bool applying_loras_ = false;
SDVersion lora_version_ = VERSION_COUNT;
uint64_t current_lora_epoch_ = 0;
uint64_t residency_epoch_ = 0;
@ -102,6 +111,7 @@ private:
void finish_compute_backend_usage(const std::vector<TensorState*>& states);
void release_all();
void invalidate_sources(const std::unordered_set<TensorState*>& states);
ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend);
bool populate_prefetch_block(PrefetchBlock& block);
@ -152,15 +162,27 @@ private:
void free_params_storage_block(ParamsStorageBlock& block);
void erase_params_storage_block(ParamsStorageBlock* block);
void reset_lora_applied_params();
bool unregister_tensor_states(const std::unordered_set<TensorState*>& states, size_t* size);
size_t other_runtime_resident_bytes(uintptr_t owner_id,
ggml_backend_t compute_backend) const;
public:
~ModelManager() override;
ModelLoader& loader() { return model_loader_; }
const ModelLoader& loader() const { return model_loader_; }
bool set_loader(ModelLoader loader);
bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false);
bool del_file(ModelLoader::FileId id);
bool refresh_files();
ModelLoader::FileVersions source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const;
size_t registered_params_size(const std::set<ModelComponent>& components) const;
void prepare_file_io() { model_loader_.process_model_files(enable_mmap_, writable_mmap_); }
bool load_float_tensor(const std::string& name, std::vector<float>& data) {
return model_loader_.load_float_tensor(name, data, n_threads_, enable_mmap_);
}
void set_n_threads(int n_threads) {
n_threads_ = n_threads;
model_loader_.set_n_threads(n_threads);
@ -172,14 +194,15 @@ public:
void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
void set_loras(std::vector<LoraSpec> loras, SDVersion version);
bool prepare_lora_sources(std::vector<LoraSpec>& loras);
bool set_loras(std::vector<LoraSpec> loras, SDVersion version);
void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector<std::pair<ggml_backend_t, size_t>>& device_limits);
static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
std::set<std::string> tensor_names() const;
bool register_param_tensors(const std::string& desc,
bool register_param_tensors(ModelComponent component,
std::map<std::string, ggml_tensor*> tensors,
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
@ -187,13 +210,18 @@ public:
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr,
ModelLoader::FileId source_file = 0,
SDVersion source_version = VERSION_COUNT);
bool unregister_param_tensors(const std::string& desc,
ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const override;
bool unregister_param_tensors(const std::vector<ggml_tensor*>& tensors);
bool unregister_param_tensors(ModelComponent component,
size_t* registered_tensor_size = nullptr);
template <typename Runner>
bool register_runner_params(const std::string& desc,
bool register_runner_params(ModelComponent component,
Runner& runner,
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
@ -201,7 +229,7 @@ public:
size_t* registered_tensor_size = nullptr) {
std::map<std::string, ggml_tensor*> tensors;
runner.get_param_tensors(tensors);
return register_param_tensors(desc,
return register_param_tensors(component,
std::move(tensors),
residency_mode,
compute_backend,
@ -210,7 +238,7 @@ public:
}
template <typename Runner>
bool register_runner_params(const std::string& desc,
bool register_runner_params(ModelComponent component,
Runner& runner,
const std::string& prefix,
ResidencyMode residency_mode,
@ -219,7 +247,7 @@ public:
size_t* registered_tensor_size = nullptr) {
std::map<std::string, ggml_tensor*> tensors;
runner.get_param_tensors(tensors, prefix);
return register_param_tensors(desc,
return register_param_tensors(component,
std::move(tensors),
residency_mode,
compute_backend,

154
src/model_manager_files.cpp Normal file
View File

@ -0,0 +1,154 @@
#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;
}
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 = sources_for(*state);
auto source = sources.find(state->name);
const bool found = source != sources.end();
if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
(lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
changed.insert(state.get());
}
}
invalidate_sources(changed);
for (auto* state : changed) {
const auto& sources = sources_for(*state);
auto source = sources.find(state->name);
state->has_source = source != sources.end();
state->source = state->has_source ? source->second : TensorStorage{};
}
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;
}
auto source = sources.find(state->name);
if (source != sources.end()) {
versions[source->second.file_id] = source->second.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;
}

1733
src/stable-diffusion.cpp
View File

File diff suppressed because it is too large Load Diff

5
src/upscaler.cpp
View File

@ -73,7 +73,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
model_manager->set_n_threads(n_threads);
model_manager->set_enable_mmap(false);
ModelLoader& model_loader = model_manager->loader();
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(esrgan_path, "", VERSION_ESRGAN)) {
LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
return false;
@ -94,7 +94,8 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
std::map<std::string, ggml_tensor*> tensors;
esrgan_upscaler->get_param_tensors(tensors);
if (!model_manager->register_param_tensors("ESRGAN",
if (!model_manager->set_loader(model_loader) ||
!model_manager->register_param_tensors(ModelComponent::Upscaler,
std::move(tensors),
backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
backend_for(SDBackendModule::UPSCALER),