From 587e91405610c7d44a04432e6e786d004404e1fe Mon Sep 17 00:00:00 2001 From: leejet Date: Thu, 10 Sep 2026 23:36:00 +0800 Subject: [PATCH] refactor: unify model source and weight lifecycle management --- src/core/ggml_runner.cpp | 27 +- src/core/ordered_map.hpp | 18 +- src/detailer.cpp | 5 +- src/device_residency_manager.h | 7 +- src/extensions/generation_extension.h | 5 +- src/extensions/photomaker_extension.cpp | 14 +- src/extensions/pulid_extension.cpp | 4 +- src/model/adapter/lora.hpp | 179 +- src/model/adapter/pmid.hpp | 92 -- src/model/diffusion/control.hpp | 40 +- src/model/diffusion/flux.hpp | 7 +- src/model/diffusion/ltxv.hpp | 7 +- src/model/diffusion/mmdit.hpp | 5 +- src/model/diffusion/qwen_image.hpp | 7 +- src/model/diffusion/wan.hpp | 7 +- src/model/diffusion/z_image.hpp | 7 +- src/model/te/llm.hpp | 7 +- src/model/te/t5.hpp | 7 +- src/model/upscaler/ltx_latent_upscaler.hpp | 9 +- src/model/vae/ltx_audio_vae.hpp | 7 +- src/model/vae/ltx_vae.hpp | 7 +- src/model/vae/wan_vae.hpp | 5 +- src/model_builders.cpp | 574 +++++++ src/model_builders.h | 63 + src/model_component.h | 61 + src/model_io/safetensors_io.cpp | 39 +- src/model_io/safetensors_io.h | 4 +- src/model_io/tensor_storage.h | 8 +- src/model_loader.cpp | 113 +- src/model_loader.h | 66 +- src/model_loader_files.cpp | 336 ++++ src/model_manager.cpp | 327 ++-- src/model_manager.h | 54 +- src/model_manager_files.cpp | 154 ++ src/stable-diffusion.cpp | 1727 +++++++++----------- src/upscaler.cpp | 5 +- 36 files changed, 2545 insertions(+), 1459 deletions(-) create mode 100644 src/model_builders.cpp create mode 100644 src/model_builders.h create mode 100644 src/model_component.h create mode 100644 src/model_loader_files.cpp create mode 100644 src/model_manager_files.cpp diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 93faa0e43..3737dc7aa 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -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(this)); std::vector 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(tensor)); + if (parameter != nullptr) + tensors.push_back(parameter); } manager->evict_compute_backend_params(tensors); manager->remove_runtime_owner(reinterpret_cast(this)); @@ -620,7 +615,15 @@ std::optional> 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) { diff --git a/src/core/ordered_map.hpp b/src/core/ordered_map.hpp index fc4755b87..308d81309 100644 --- a/src/core/ordered_map.hpp +++ b/src/core/ordered_map.hpp @@ -38,9 +38,19 @@ class OrderedMap { 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 @@ class OrderedMap { } }; -#endif // __SD_CORE_ORDERED_MAP_HPP__ \ No newline at end of file +#endif // __SD_CORE_ORDERED_MAP_HPP__ diff --git a/src/detailer.cpp b/src/detailer.cpp index 2523dd513..bbf9c6880 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -676,7 +676,7 @@ bool ADetailerGGML::load_from_file(const std::string& detector_path) { model_manager = std::make_shared(); 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 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 diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index 2fb2a342d..0a2a8895f 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -48,9 +48,10 @@ struct DeviceResidencyManager { const std::vector& required_params) const = 0; virtual bool assign_compute_backend(const std::vector& tensors, ggml_backend_t compute_backend) = 0; - virtual bool prepare_params(const std::vector& tensors) = 0; - virtual void release_compute_backend_params(const std::vector& tensors) = 0; - virtual void evict_compute_backend_params(const std::vector& tensors) = 0; + virtual ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const { return nullptr; } + virtual bool prepare_params(const std::vector& tensors) = 0; + virtual void release_compute_backend_params(const std::vector& tensors) = 0; + virtual void evict_compute_backend_params(const std::vector& tensors) = 0; virtual WeightResidencyInfo inspect_compute_backend_params( const std::vector& tensors) const = 0; virtual void update_runtime_residency(uintptr_t owner_id, diff --git a/src/extensions/generation_extension.h b/src/extensions/generation_extension.h index b1ed9729b..e5ffd770e 100644 --- a/src/extensions/generation_extension.h +++ b/src/extensions/generation_extension.h @@ -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 model_manager; int n_threads; std::function 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; } diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp index b3bd6ec26..05bb46c0d 100644 --- a/src/extensions/photomaker_extension.cpp +++ b/src/extensions/photomaker_extension.cpp @@ -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(ctx.backend_for(SDBackendModule::PHOTOMAKER), + pmid_model = std::make_shared(ctx.backend_for(SDBackendModule::PHOTOMAKER), ctx.tensor_storage_map, "pmid", ctx.version, diff --git a/src/extensions/pulid_extension.cpp b/src/extensions/pulid_extension.cpp index d529e5710..fa660c300 100644 --- a/src/extensions/pulid_extension.cpp +++ b/src/extensions/pulid_extension.cpp @@ -79,8 +79,8 @@ struct PuLIDExtension : public GenerationExtension { sd::Tensor 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 { diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 07be8a768..738414314 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -23,25 +23,31 @@ struct LoraModel : public GGMLRunner { std::set skipped_incompatible_lora_tensors; std::set warned_incompatible_model_tensors; std::string file_path; - std::shared_ptr 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 registered_params; + std::map scalar_values; typedef std::function 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 manager = std::make_shared()) - : 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 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(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(residency_manager.lock()); + if (model_manager == nullptr) return false; - } - - std::unordered_map 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 lock(lora_mutex); - tensors_to_create[name] = tensor_storage; - } + 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 tensors; + std::map scalars; + std::set 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 { - const std::string& name = tensor_storage.name; - auto iter = lora_tensors.find(name); - if (iter != lora_tensors.end()) { - *dst_tensor = iter->second; - } + tensors[name] = tensor; } - 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; } - - std::map tensors; - for (const auto& pair : lora_tensors) { - tensors[pair.first] = pair.second; - } - 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"); - return false; + // 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; } - std::vector 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"); + 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; - } + 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(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(); - 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 model_tensors, + bool apply(std::map model_tensors, const std::set& 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 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 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) { diff --git a/src/model/adapter/pmid.hpp b/src/model/adapter/pmid.hpp index f16d7ba1f..900db5cb1 100644 --- a/src/model/adapter/pmid.hpp +++ b/src/model/adapter/pmid.hpp @@ -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 @@ struct PhotoMakerIDEncoder : public GGMLRunner { } }; -struct PhotoMakerIDEmbed : public GGMLRunner { - std::map tensors; - std::string file_path; - std::shared_ptr 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 manager = std::make_shared(), - 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 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 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::iterator pos; - pos = tensors.find("pmid.id_embeds"); - if (pos != tensors.end()) - return pos->second; - return nullptr; - } -}; - #endif // __SD_MODEL_ADAPTER_PMID_HPP__ diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index 66cce5759..34387d1b7 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -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> controls; bool guided_hint_cached = false; - std::shared_ptr 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 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 tensors; - control_net.get_param_tensors(tensors); - - auto manager = std::dynamic_pointer_cast(residency_manager.lock()); - if (manager == nullptr) { - owned_model_manager = std::make_shared(); - 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__ diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 0ac54a9a1..67e79d3d9 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index cddfc2d3b..118606e4c 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 9063b65d1..454ef4b3e 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -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, diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index a9dfdb838..07364958f 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 1884be5c7..687f28380 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index b29e75457..4ae47e268 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index d381ac5bb..4f1db3bca 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 4c86117cc..f89488d36 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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 t5 = std::make_shared(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, diff --git a/src/model/upscaler/ltx_latent_upscaler.hpp b/src/model/upscaler/ltx_latent_upscaler.hpp index 88c6b74ad..6650adc70 100644 --- a/src/model/upscaler/ltx_latent_upscaler.hpp +++ b/src/model/upscaler/ltx_latent_upscaler.hpp @@ -433,12 +433,15 @@ namespace LTXVUpsampler { struct LatentUpsamplerRunner : public GGMLRunner { LatentUpsamplerConfig config; std::unique_ptr model; + std::string weight_prefix; LatentUpsamplerRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "", std::shared_ptr 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(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& tensors) { if (model) { - model->get_param_tensors(tensors); + model->get_param_tensors(tensors, weight_prefix); } } diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 3c9955adf..97c434146 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index 93f6b760e..ac22c09b4 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -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(); - ModelLoader& model_loader = model_manager->loader(); + auto model_manager = std::make_shared(); + 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, diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index bfdb4b7d5..a632d66b4 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -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, diff --git a/src/model_builders.cpp b/src/model_builders.cpp new file mode 100644 index 000000000..8d174e8d7 --- /dev/null +++ b/src/model_builders.cpp @@ -0,0 +1,574 @@ +#include "model_builders.h" + +#include +#include + +#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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager); + result.diffusion = std::make_shared(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(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + } else { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + } + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "text_embedding_projection", + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + true, + weight_manager); + result.diffusion = std::make_shared(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(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(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model", + weight_manager); + } else if (sd_version_is_minit2i(version)) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + } else { // SD1.x SD2.x SDXL + std::map 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(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + embbeding_map, + version, + weight_manager); + result.diffusion = std::make_shared(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(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(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 { + 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(ctx.backends.runtime_backend(SDBackendModule::VAE), + tensor_storage_map, + "decoder", + decode_only, + version, + weight_manager); + + } else { + auto model = std::make_shared(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 { + if (sd_version_is_ltxav(version)) { + return std::make_shared(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(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(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(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(ctx.backends.runtime_backend(SDBackendModule::VAE), + tensor_storage_map, + "first_stage_model", + false, + vae_version, + weight_manager); + } else { + auto model = std::make_shared(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(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(ctx.backends.runtime_backend(SDBackendModule::VAE), + tensor_storage_map, + "", + weight_manager); + } else { + result.audio = std::make_shared(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& 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(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>& extensions) { + std::vector> 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 diff --git a/src/model_builders.h b/src/model_builders.h new file mode 100644 index 000000000..ac95a6f65 --- /dev/null +++ b/src/model_builders.h @@ -0,0 +1,63 @@ +#ifndef __SD_MODEL_BUILDERS_H__ +#define __SD_MODEL_BUILDERS_H__ + +#include +#include + +#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 weight_manager; + }; + + struct CoreRunners { + std::shared_ptr conditioner; + std::shared_ptr diffusion; + std::shared_ptr high_noise_diffusion; + std::shared_ptr clip_vision; + std::shared_ptr 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; + std::shared_ptr preview; + std::shared_ptr 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& runner); + bool build_extension_runners(const GenerationExtensionInitContext& ctx, + std::vector>& extensions); + +} // namespace sd::model_builders + +#endif // __SD_MODEL_BUILDERS_H__ diff --git a/src/model_component.h b/src/model_component.h new file mode 100644 index 000000000..96077a4a4 --- /dev/null +++ b/src/model_component.h @@ -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__ diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index dfcb4a929..807c915aa 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -144,7 +144,8 @@ static bool read_comfy_quant_config(std::ifstream& file, bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, std::string* error, - std::map* metadata) { + std::map* metadata, + std::set* 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 end = tensor_info["data_offsets"][1].get(); 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(); + if (size < 0 || (bytes != 0 && static_cast(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(); + if (ne[i] < 0 || (elements != 0 && static_cast(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) { diff --git a/src/model_io/safetensors_io.h b/src/model_io/safetensors_io.h index 4291b54f6..846352729 100644 --- a/src/model_io/safetensors_io.h +++ b/src/model_io/safetensors_io.h @@ -2,6 +2,7 @@ #define __SD_MODEL_IO_SAFETENSORS_IO_H__ #include +#include #include #include @@ -12,7 +13,8 @@ bool is_safetensors_file(const std::string& file_path); bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, std::string* error = nullptr, - std::map* metadata = nullptr); + std::map* metadata = nullptr, + std::set* tensor_names = nullptr); bool read_safetensors_index_file(const std::string& file_path, std::vector& shard_paths, std::string* error = nullptr); diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 11f53e063..6d2562849 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -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; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 9252240b6..8325a6063 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -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(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 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())) + .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 ModelLoader::get_wtype_stat() { +std::map ModelLoader::get_wtype_stat() const { std::map wtype_stat; for (auto& [name, tensor_storage] : tensor_storage_map) { if (is_unused_tensor(tensor_storage.name)) { @@ -640,7 +666,7 @@ std::map ModelLoader::get_wtype_stat() { return wtype_stat; } -std::map ModelLoader::get_conditioner_wtype_stat() { +std::map ModelLoader::get_conditioner_wtype_stat() const { std::map wtype_stat; for (auto& [name, tensor_storage] : tensor_storage_map) { if (is_unused_tensor(tensor_storage.name)) { @@ -664,7 +690,7 @@ std::map ModelLoader::get_conditioner_wtype_stat() { return wtype_stat; } -std::map ModelLoader::get_diffusion_model_wtype_stat() { +std::map ModelLoader::get_diffusion_model_wtype_stat() const { std::map wtype_stat; for (auto& [name, tensor_storage] : tensor_storage_map) { if (is_unused_tensor(tensor_storage.name)) { @@ -685,7 +711,7 @@ std::map ModelLoader::get_diffusion_model_wtype_stat() { return wtype_stat; } -std::map ModelLoader::get_vae_wtype_stat() { +std::map ModelLoader::get_vae_wtype_stat() const { std::map 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 ModelLoader::mmap_tensors(std::map& tensors, std::set ignore_tensors, bool writable_mmap) { + std::set names; + for (const auto& entry : tensors) { + names.insert(entry.first); + } + if (!validate_sources(&names)) { + return {}; + } process_model_files(true, writable_mmap); std::vector result; @@ -919,6 +957,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool enable_mmap, const std::set* target_tensor_names, bool log_progress) { + if (!validate_sources(target_tensor_names)) { + return false; + } process_model_files(enable_mmap, false); std::atomic 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& 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); diff --git a/src/model_loader.h b/src/model_loader.h index f7ebcf3ef..c5866b263 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -2,6 +2,7 @@ #define __MODEL_LOADER_H__ #include +#include #include #include #include @@ -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; + 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 dependencies; + String2TensorStorage tensors; + std::map metadata; + }; + + std::vector files_; + uint64_t revision_ = 0; + bool names_converted_ = false; + ggml_type wtype_override_ = GGML_TYPE_COUNT; + std::string tensor_type_rules_; + std::vector parsed_dependencies_; + std::map> 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 file_paths_; @@ -52,16 +93,27 @@ class ModelLoader { 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& 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* tensor_names = nullptr) const; + uint64_t revision() const { return revision_; } + FileVersions file_versions(const std::vector& 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 get_wtype_stat(); - std::map get_conditioner_wtype_stat(); - std::map get_diffusion_model_wtype_stat(); - std::map get_vae_wtype_stat(); + SDVersion get_sd_version() const; + std::map get_wtype_stat() const; + std::map get_conditioner_wtype_stat() const; + std::map get_diffusion_model_wtype_stat() const; + std::map 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& get_metadata() const { return metadata_; } @@ -92,8 +144,8 @@ class ModelLoader { 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; }; diff --git a/src/model_loader_files.cpp b/src/model_loader_files.cpp new file mode 100644 index 000000000..181b619a4 --- /dev/null +++ b/src/model_loader_files.cpp @@ -0,0 +1,336 @@ +#include "model_loader.h" + +#include +#include +#include + +#include "core/util.h" +#include "name_conversion.h" + +static uint64_t next_source_revision() { + static std::atomic 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 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 file_indices; + std::vector 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(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 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* tensor_names) const { + std::set 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& 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& 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); +} diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 166030a5c..f3bf19a2c 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include "core/ggml_extend_backend.h" @@ -26,7 +27,8 @@ static bool lora_specs_equal(const std::vector& 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 ignore_tensor common_ignore_tensors_ = std::move(ignore_tensors); } -void ModelManager::set_loras(std::vector loras, SDVersion version) { - if (loras.empty() && loras_.empty()) { - lora_version_ = version; - return; +bool ModelManager::prepare_lora_sources(std::vector& loras) { + ModelLoader candidate = model_loader_; + std::vector resolved; + std::set 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 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 ModelManager::tensor_names() const { std::set 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 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* tensor_ops) { - if (desc.empty()) { - LOG_ERROR("model manager tensor desc is empty"); + const std::map* 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 new_tensors; std::vector> 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(); - state->name = name; - state->tensor = tensor; - state->desc = desc; + auto state = std::make_unique(); + 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 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 target_states; +bool ModelManager::unregister_param_tensors(const std::vector& tensors) { + std::unordered_set 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& 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 storage_blocks_to_release; std::unordered_set 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& 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 model_tensors; std::vector states; }; - std::map groups; + using ApplyTarget = std::tuple; + std::map 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& 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& states std::set 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(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(this, [](ModelManager*) {}); + auto lora = std::make_shared(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& 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& 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& states } void ModelManager::reset_lora_applied_params() { - clear_all_prefetched_params(); - release_compute_staging_blocks(true); - release_params_storage_blocks(true); + std::unordered_set 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& 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,76 +926,56 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, } bool ModelManager::load_tensors(const std::vector& states) { - std::map states_by_name; - std::set target_tensor_names; - for (TensorState* state : states) { - if (state == nullptr) { + using ReadGroup = std::pair; + using ReadBatch = std::map>; + std::map> groups; + for (auto* state : states) { + if (state == nullptr) continue; - } - states_by_name[state->name] = state; - target_tensor_names.insert(state->name); - } - if (states_by_name.empty()) { - return true; - } - - std::set 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; - } - - 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 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; + 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()); + } + (*batch)[state->name].push_back(state); + } + for (auto& group : groups) { + for (auto& batch : group.second) { + std::set names; + std::set 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 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; + } + } } } - 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& tensors, std::vector& required_states, ggml_backend_t compute_backend) const { @@ -1147,21 +1204,13 @@ bool ModelManager::resolve_required_tensor_states(const std::vectorsecond; - 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& 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 failed(required_states.begin(), required_states.end()); + invalidate_sources(failed); return false; } return true; diff --git a/src/model_manager.h b/src/model_manager.h index b23f0db63..ccd7a5e34 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -10,6 +10,7 @@ #include #include "device_residency_manager.h" +#include "model_component.h" #include "model_loader.h" class ModelManager : public DeviceResidencyManager { @@ -24,7 +25,9 @@ class ModelManager : public DeviceResidencyManager { 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 @@ class ModelManager : public DeviceResidencyManager { 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 @@ class ModelManager : public DeviceResidencyManager { ModelLoader model_loader_; std::vector> tensor_states_; - std::map tensor_states_by_name_; + std::map tensor_states_by_tensor_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; @@ -91,6 +98,8 @@ class ModelManager : public DeviceResidencyManager { bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; + std::set 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 @@ class ModelManager : public DeviceResidencyManager { void finish_compute_backend_usage(const std::vector& states); void release_all(); + void invalidate_sources(const std::unordered_set& states); ggml_backend_t prefetch_backend_for(ggml_backend_t compute_backend); bool populate_prefetch_block(PrefetchBlock& block); @@ -152,15 +162,27 @@ class ModelManager : public DeviceResidencyManager { 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& 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& components, const ModelLoader& loader) const; + size_t registered_params_size(const std::set& 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& 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 @@ class ModelManager : public DeviceResidencyManager { 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 ignore_tensors); - void set_loras(std::vector loras, SDVersion version); + bool prepare_lora_sources(std::vector& loras); + bool set_loras(std::vector loras, SDVersion version); void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits); static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor); std::set tensor_names() const; - bool register_param_tensors(const std::string& desc, + bool register_param_tensors(ModelComponent component, std::map tensors, ResidencyMode residency_mode, ggml_backend_t compute_backend, @@ -187,13 +210,18 @@ class ModelManager : public DeviceResidencyManager { size_t* registered_tensor_size = nullptr, bool allow_split_buffer = false, bool params_follow_compute_backend = false, - const std::map* tensor_ops = nullptr); + const std::map* tensor_ops = nullptr, + ModelLoader::FileId source_file = 0, + SDVersion source_version = VERSION_COUNT); + + ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const override; + bool unregister_param_tensors(const std::vector& tensors); - bool unregister_param_tensors(const std::string& desc, + bool unregister_param_tensors(ModelComponent component, size_t* registered_tensor_size = nullptr); template - 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 @@ class ModelManager : public DeviceResidencyManager { size_t* registered_tensor_size = nullptr) { std::map 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 @@ class ModelManager : public DeviceResidencyManager { } template - 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 @@ class ModelManager : public DeviceResidencyManager { size_t* registered_tensor_size = nullptr) { std::map tensors; runner.get_param_tensors(tensors, prefix); - return register_param_tensors(desc, + return register_param_tensors(component, std::move(tensors), residency_mode, compute_backend, diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp new file mode 100644 index 000000000..b5dfe7c7c --- /dev/null +++ b/src/model_manager_files.cpp @@ -0,0 +1,154 @@ +#include "model_manager.h" + +#include + +#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& 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, 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 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& 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& components) const { + size_t bytes = 0; + std::unordered_set 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; +} diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index dfc961ebc..442074e80 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2,6 +2,8 @@ #include #include #include +#include +#include #include #include #include @@ -19,6 +21,7 @@ #include "core/rng_mt19937.hpp" #include "core/rng_philox.hpp" #include "core/util.h" +#include "model_builders.h" #include "model_loader.h" #include "model_manager.h" #include "stable-diffusion.h" @@ -28,41 +31,15 @@ #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" #include "model/adapter/lora.hpp" -#include "model/diffusion/anima.hpp" #include "model/diffusion/animatediff.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/upscaler/esrgan.hpp" #include "model/upscaler/ltx_latent_upscaler.hpp" -#include "model/vae/auto_encoder_kl.hpp" -#include "model/vae/hunyuan_vae.hpp" -#include "model/vae/ltx_audio_vae.hpp" +#include "model/vae/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" #include "runtime/denoiser.hpp" #include "runtime/guidance.h" #include "runtime/preview_interval.h" @@ -75,7 +52,6 @@ #include const char* sd_vae_format_name(enum sd_vae_format_t format); -static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback); static bool sd_version_supports_animatediff(SDVersion version) { return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX; @@ -243,9 +219,18 @@ class StableDiffusionGGML { sd::Tensor ip_adapter_uncond_tokens; float ip_adapter_strength = 1.0f; std::vector> generation_extensions; - std::vector> runtime_lora_models; + struct RuntimeLora { + ModelManager::LoraSpec spec; + SDBackendModule module; + std::shared_ptr model; + + bool matches(const ModelManager::LoraSpec& other) const { + return spec.file_id == other.file_id && spec.file_revision == other.file_revision && + spec.tensor_name_prefix_filter == other.tensor_name_prefix_filter; + } + }; + std::vector runtime_lora_models; bool apply_lora_immediately = false; - bool animatediff_loaded = false; int animatediff_num_frames = 0; std::string taesd_path; @@ -265,17 +250,275 @@ class StableDiffusionGGML { bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; - size_t control_net_params_mem_size = 0; - std::shared_ptr model_manager; + enum class RunnerGroup { Core, + VAE, + ControlNet, + Extensions }; + using RunnerGroups = std::set; + + struct ModelConfig { + sd_ctx_params_t params{}; + std::list strings; + std::vector embeddings; + ModelLoader::FileId control_net_file = 0; + bool use_tae = false; + bool use_audio_vae = false; + bool photomaker_source_available = false; + bool animatediff_loaded = false; + + explicit ModelConfig(const sd_ctx_params_t& initial) + : params(initial) { + for (auto member : {&sd_ctx_params_t::model_path, &sd_ctx_params_t::clip_l_path, + &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path, + &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path, + &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path, + &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path, + &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path, + &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path, + &sd_ctx_params_t::control_net_path, &sd_ctx_params_t::ip_adapter_path, + &sd_ctx_params_t::motion_module_path, &sd_ctx_params_t::photo_maker_path, + &sd_ctx_params_t::pulid_weights_path, &sd_ctx_params_t::tensor_type_rules, + &sd_ctx_params_t::max_vram, &sd_ctx_params_t::backend, + &sd_ctx_params_t::params_backend, &sd_ctx_params_t::split_mode, + &sd_ctx_params_t::rpc_servers, &sd_ctx_params_t::model_args}) { + strings.emplace_back(SAFE_STR(initial.*member)); + params.*member = strings.back().c_str(); + } + for (uint32_t i = 0; i < initial.embedding_count; ++i) { + strings.emplace_back(SAFE_STR(initial.embeddings[i].name)); + const char* name = strings.back().c_str(); + strings.emplace_back(SAFE_STR(initial.embeddings[i].path)); + embeddings.push_back({name, strings.back().c_str()}); + } + params.embeddings = embeddings.data(); + } + + ModelConfig(const ModelConfig& other) + : ModelConfig(other.params) { + control_net_file = other.control_net_file; + use_tae = other.use_tae; + use_audio_vae = other.use_audio_vae; + photomaker_source_available = other.photomaker_source_available; + animatediff_loaded = other.animatediff_loaded; + } + ModelConfig& operator=(const ModelConfig&) = delete; + + void set_control_net(ModelLoader::FileId id, const std::string& path) { + control_net_file = id; + strings.push_back(path); + params.control_net_path = strings.back().c_str(); + } + }; + + struct RunnerState { + bool ready = false; + uint64_t catalog_revision = 0; + std::map sources; + }; + + std::recursive_mutex execution_mutex; + std::unique_ptr config_; + RunnerState runner_state_; + bool executing_ = false; + std::shared_ptr denoiser = std::make_shared(); std::vector file_alphas_cumprod; - StableDiffusionGGML() = default; - + StableDiffusionGGML() = default; ~StableDiffusionGGML() = default; + static const std::map>& runner_components() { + static const std::map> components{ + {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}}, + {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}}, + {RunnerGroup::ControlNet, {ModelComponent::ControlNet}}, + {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}}, + }; + return components; + } + + static RunnerGroups all_runner_groups() { + RunnerGroups groups; + for (const auto& entry : runner_components()) { + groups.insert(entry.first); + } + return groups; + } + + ModelLoader::FileVersions runner_source_versions(RunnerGroup group, const ModelLoader& loader) const { + auto sources = model_manager->source_versions(runner_components().at(group), loader); + if (group == RunnerGroup::Core) { + // PhotoMaker's LoRA may already be merged into resident core weights. + auto extra = loader.file_versions({"alphas_cumprod", "v_pred", "edm_vpred.", "pmid."}); + sources.insert(extra.begin(), extra.end()); + } + return sources; + } + + void capture_runner_sources() { + RunnerState state; + state.catalog_revision = model_manager->loader().revision(); + for (const auto& entry : runner_components()) { + state.sources[entry.first] = runner_source_versions(entry.first, model_manager->loader()); + } + state.ready = true; + runner_state_ = std::move(state); + } + + void end_runners() { + if (cond_stage_model) + cond_stage_model->runner_end(); + if (diffusion_model) + diffusion_model->runner_end(); + if (high_noise_diffusion_model) + high_noise_diffusion_model->runner_end(); + if (clip_vision) + clip_vision->runner_end(); + if (ip_adapter) + ip_adapter->runner_end(); + if (first_stage_model) + first_stage_model->runner_end(); + if (preview_vae) + preview_vae->runner_end(); + if (audio_vae_model) + audio_vae_model->runner_end(); + if (control_net) + control_net->runner_end(); + for (auto& extension : generation_extensions) + extension->runner_end(); + for (auto& lora : runtime_lora_models) + if (lora.model) + lora.model->runner_end(); + } + + bool reset_runners(const RunnerGroups& groups) { + end_runners(); + clear_lora_adapters(); + runtime_lora_models.clear(); + for (auto group : groups) { + for (auto component : runner_components().at(group)) { + if (!model_manager->unregister_param_tensors(component)) { + return false; + } + } + } + for (auto group : groups) { + switch (group) { + case RunnerGroup::Core: + cond_stage_model.reset(); + diffusion_model.reset(); + high_noise_diffusion_model.reset(); + clip_vision.reset(); + ip_adapter.reset(); + ip_adapter_tokens = {}; + ip_adapter_uncond_tokens = {}; + runtime_lora_models.clear(); + break; + case RunnerGroup::VAE: + first_stage_model.reset(); + preview_vae.reset(); + audio_vae_model.reset(); + break; + case RunnerGroup::ControlNet: + control_net.reset(); + break; + case RunnerGroup::Extensions: + generation_extensions.clear(); + break; + } + } + return true; + } + + bool refresh_model_sources() { + bool changed; + if (!model_manager->loader().files_changed(changed, false)) { + return false; + } + if (!changed && runner_state_.ready && runner_state_.catalog_revision == model_manager->loader().revision()) { + return true; + } + ModelLoader candidate = model_manager->loader(); + return candidate.refresh_files(false) && apply_model_update(std::move(candidate)); + } + + bool apply_model_update(ModelLoader candidate, + std::unique_ptr next_config = nullptr, + RunnerGroups groups = {}) { + const SDVersion next_version = candidate.get_sd_version(); + if (next_version == VERSION_COUNT) { + LOG_ERROR("cannot identify updated diffusion model"); + return false; + } + if (!runner_state_.ready || next_version != version) { + groups = all_runner_groups(); + } else { + for (const auto& entry : runner_components()) { + if (runner_state_.sources.at(entry.first) != runner_source_versions(entry.first, candidate)) { + groups.insert(entry.first); + } + } + } + runner_state_.ready = false; + if (!reset_runners(groups)) { + return false; + } + if (!model_manager->set_loader(std::move(candidate))) { + reset_runners(all_runner_groups()); + return false; + } + if (next_config) { + config_ = std::move(next_config); + } + version = next_version; + if (!build_runners(groups)) { + reset_runners(all_runner_groups()); + return false; + } + capture_runner_sources(); + return true; + } + + struct ContextOperation { + StableDiffusionGGML& sd; + std::unique_lock lock; + bool acquired = false; + + explicit ContextOperation(StableDiffusionGGML& sd) + : sd(sd), lock(sd.execution_mutex, std::try_to_lock) { + if (!lock.owns_lock() || sd.executing_) { + // The caller may be a log callback, so rejecting it must not log. + return; + } + sd.executing_ = true; + acquired = true; + } + + ~ContextOperation() { + if (acquired) { + sd.executing_ = false; + } + } + }; + + struct ExecutionScope { + ContextOperation operation; + bool ready = false; + + explicit ExecutionScope(StableDiffusionGGML& sd) + : operation(sd) { + ready = operation.acquired && sd.refresh_model_sources(); + } + + ~ExecutionScope() { + if (ready) { + operation.sd.end_runners(); + } + } + }; + ggml_backend_t backend_for(SDBackendModule module) { ggml_backend_t module_backend = backend_manager.runtime_backend(module); if (module_backend == nullptr) { @@ -327,7 +570,7 @@ class StableDiffusionGGML { } template - bool register_runner_params(const std::string& desc, + bool register_runner_params(ModelComponent component, const std::shared_ptr& model, SDBackendModule module, size_t* params_mem_size = nullptr) { @@ -351,7 +594,7 @@ class StableDiffusionGGML { if constexpr (has_set_runtime_backends::value) { if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) { if (backend_manager.split_mode(module) == SDSplitMode::ROW) { - return register_row_split_runner_params(desc, + return register_row_split_runner_params(component, model, module, module_backends, @@ -360,7 +603,7 @@ class StableDiffusionGGML { residency_mode, params_mem_size); } - return register_layer_split_runner_params(desc, + return register_layer_split_runner_params(component, model, module, module_backends, @@ -374,7 +617,7 @@ class StableDiffusionGGML { sd_backend_module_name(module), sd::layer_split_backend_device_display_name(module_backends[0]).c_str()); } - return model_manager->register_param_tensors(desc, + return model_manager->register_param_tensors(component, std::move(group_tensors), residency_mode, backend_for(module), @@ -386,7 +629,7 @@ class StableDiffusionGGML { } template - bool register_row_split_runner_params(const std::string& desc, + bool register_row_split_runner_params(ModelComponent component, const std::shared_ptr& model, SDBackendModule module, const std::vector& module_backends, @@ -397,8 +640,8 @@ class StableDiffusionGGML { ggml_backend_t main_backend = module_backends[0]; auto fall_back_to_layer_split = [&](const char* reason) { - LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason); - return register_layer_split_runner_params(desc, + LOG_WARN("%s: row split unavailable (%s); falling back to layer split", model_component_name(component), reason); + return register_layer_split_runner_params(component, model, module, module_backends, @@ -470,13 +713,13 @@ class StableDiffusionGGML { } LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)", - desc.c_str(), + model_component_name(component), row_split_map.size(), row_split_bytes / (1024.f * 1024.f), module_backends.size(), sd::layer_split_backend_device_display_name(main_backend).c_str()); - if (!model_manager->register_param_tensors(desc, + if (!model_manager->register_param_tensors(component, std::move(row_split_map), residency_mode, main_backend, @@ -487,7 +730,7 @@ class StableDiffusionGGML { &tensor_ops)) { return false; } - return model_manager->register_param_tensors(desc, + return model_manager->register_param_tensors(component, std::move(regular_map), residency_mode, main_backend, @@ -502,7 +745,7 @@ class StableDiffusionGGML { // The first real graph assigns each param tensor to a runtime backend // before weights are loaded or staged. template - bool register_layer_split_runner_params(const std::string& desc, + bool register_layer_split_runner_params(ModelComponent component, const std::shared_ptr& model, SDBackendModule module, const std::vector& module_backends, @@ -521,10 +764,10 @@ class StableDiffusionGGML { LOG_WARN( "%s: layer split across a CPU device is not supported; using %s " "(use --params-backend %s=cpu to keep weights in RAM)", - desc.c_str(), + model_component_name(component), sd::layer_split_backend_device_display_name(module_backends[0]).c_str(), sd_backend_module_name(module)); - return model_manager->register_param_tensors(desc, + return model_manager->register_param_tensors(component, std::move(group_tensors), residency_mode, module_backends[0], @@ -546,11 +789,11 @@ class StableDiffusionGGML { } LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph", - desc.c_str(), + model_component_name(component), group_tensors.size(), module_backends.size()); - return model_manager->register_param_tensors(desc, + return model_manager->register_param_tensors(component, std::move(group_tensors), residency_mode, module_backends[0], @@ -562,75 +805,43 @@ class StableDiffusionGGML { } bool unload_control_net() { - if (control_net == nullptr) { - return true; - } - if (model_manager != nullptr) { - if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) { - return false; - } - } - control_net.reset(); - control_net_params_mem_size = 0; - return true; - } - - bool load_control_net_from_file(const std::string& path) { - if (path.empty()) { - LOG_ERROR("sd_ctx_load_control_net: empty path"); + ContextOperation operation(*this); + if (!operation.acquired) { return false; } - if (model_manager == nullptr) { - LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized"); + if (model_manager == nullptr || config_ == nullptr) { + LOG_ERROR("cannot unload ControlNet: context is not initialized"); return false; } - - if (!unload_control_net()) { + ModelLoader candidate = model_manager->loader(); + if (config_->control_net_file != 0 && !candidate.del_file(config_->control_net_file)) { return false; } + auto next_config = std::make_unique(*config_); + next_config->set_control_net(0, ""); + return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet}); + } - ModelLoader& shared_loader = model_manager->loader(); - if (!shared_loader.init_from_file(path)) { - LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str()); + bool load_control_net_from_file(const std::string& path) { + ContextOperation operation(*this); + if (!operation.acquired) { return false; } - shared_loader.convert_tensors_name(); - - if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) { - LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable"); + if (path.empty() || model_manager == nullptr || config_ == nullptr) { + LOG_ERROR("cannot load ControlNet: invalid path or uninitialized context"); return false; } - - control_net = std::make_shared(backend_for(SDBackendModule::CONTROL_NET), - params_backend_for(SDBackendModule::CONTROL_NET), - shared_loader.get_tensor_storage_map(), - version, - "", - model_manager); - control_net->set_max_graph_vram_bytes( - max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); - if (diffusion_conv_direct) { - LOG_INFO("Using Conv2d direct in the control net"); - control_net->set_conv2d_direct_enabled(true); - } - if (!register_runner_params("ControlNet", - control_net, - SDBackendModule::CONTROL_NET, - &control_net_params_mem_size)) { - LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed"); - control_net.reset(); - control_net_params_mem_size = 0; + ModelLoader candidate = model_manager->loader(); + ModelLoader::FileId file_id; + if (!candidate.add_file(path, "", &file_id)) { return false; } - if (!model_manager->validate_registered_tensors()) { - LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed"); - unload_control_net(); + if (config_->control_net_file != 0 && config_->control_net_file != file_id && !candidate.del_file(config_->control_net_file)) { return false; } - LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)", - path.c_str(), - control_net_params_mem_size / 1024.0 / 1024.0); - return true; + auto next_config = std::make_unique(*config_); + next_config->set_control_net(file_id, path); + return apply_model_update(std::move(candidate), std::move(next_config), {RunnerGroup::ControlNet}); } bool init_backend() { @@ -692,11 +903,11 @@ class StableDiffusionGGML { } } - void load_alphas_cumprod(ModelLoader& model_loader) { + void load_alphas_cumprod() { file_alphas_cumprod.clear(); std::vector loaded_alphas; - if (!model_loader.load_float_tensor("alphas_cumprod", loaded_alphas, n_threads, enable_mmap)) { + if (!model_manager->load_float_tensor("alphas_cumprod", loaded_alphas)) { return; } if (loaded_alphas.size() != TIMESTEPS) { @@ -716,11 +927,10 @@ class StableDiffusionGGML { LOG_VERBOSE("loaded alphas_cumprod from model file"); } - bool init_model_loader(ModelLoader& model_loader, - const sd_ctx_params_t* sd_ctx_params, - bool& use_tae, - bool& use_audio_vae, - bool& use_control_net) { + bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration) { + const auto* sd_ctx_params = &configuration.params; + auto& use_tae = configuration.use_tae; + auto& use_audio_vae = configuration.use_audio_vae; if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); if (!model_loader.init_from_file(sd_ctx_params->model_path)) { @@ -838,16 +1048,14 @@ class StableDiffusionGGML { "model.diffusion_model.motion_module.")) { LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path); } else { - animatediff_loaded = true; + configuration.animatediff_loaded = true; } } if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { - if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) { + if (!model_loader.add_file(sd_ctx_params->control_net_path, "", &configuration.control_net_file)) { LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path); return false; - } else { - use_control_net = true; } } @@ -858,6 +1066,13 @@ class StableDiffusionGGML { } } + if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) { + configuration.photomaker_source_available = model_loader.add_file(sd_ctx_params->photo_maker_path, "pmid."); + if (!configuration.photomaker_source_available) { + LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path); + } + } + model_loader.convert_tensors_name(); ggml_type wtype = sd_type_to_ggml_type(sd_ctx_params->wtype); @@ -870,6 +1085,7 @@ class StableDiffusionGGML { } bool init(const sd_ctx_params_t* sd_ctx_params) { + auto configuration = std::make_unique(*sd_ctx_params); n_threads = sd_ctx_params->n_threads; enable_mmap = sd_ctx_params->enable_mmap; disable_prefetch = sd_ctx_params->disable_prefetch; @@ -891,10 +1107,6 @@ class StableDiffusionGGML { std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers); add_rpc_devices(rpc_servers_spec); - bool use_tae = false; - bool use_audio_vae = false; - bool use_control_net = false; - rng = get_rng(sd_ctx_params->rng_type); if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) { sampler_rng = get_rng(sd_ctx_params->sampler_rng_type); @@ -909,9 +1121,9 @@ class StableDiffusionGGML { model_manager->set_enable_mmap(enable_mmap); model_manager->set_segmented_compute_disabled(disable_segmented_compute); model_manager->set_prefetch_disabled(disable_prefetch); - ModelLoader& model_loader = model_manager->loader(); + ModelLoader model_loader; - if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) { + if (!init_model_loader(model_loader, *configuration)) { return false; } @@ -948,32 +1160,22 @@ class StableDiffusionGGML { eager_load = false; } - std::map wtype_stat = model_loader.get_wtype_stat(); - std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); - std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); - std::map vae_wtype_stat = model_loader.get_vae_wtype_stat(); - - auto wtype_stat_to_str = [](const std::map& m, int key_width = 8, int value_width = 5) -> std::string { - std::ostringstream oss; - bool first = true; - for (const auto& [type, count] : m) { - if (!first) - oss << "|"; - first = false; - oss << std::right << std::setw(key_width) << ggml_type_name(type) - << ": " - << std::left << std::setw(value_width) << count; - } - return oss.str(); - }; + diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; + return apply_model_update(std::move(model_loader), std::move(configuration), all_runner_groups()); + } - LOG_INFO("Weight type stat: %s", wtype_stat_to_str(wtype_stat).c_str()); - LOG_INFO("Conditioner weight type stat: %s", wtype_stat_to_str(conditioner_wtype_stat).c_str()); - LOG_INFO("Diffusion model weight type stat: %s", wtype_stat_to_str(diffusion_model_wtype_stat).c_str()); - LOG_INFO("VAE weight type stat: %s", wtype_stat_to_str(vae_wtype_stat).c_str()); + bool uses_tae() const { + return config_->use_tae || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09; + } - LOG_VERBOSE("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor)); + bool tae_preview_only() const { + return config_->params.tae_preview_only && version != VERSION_SDXS_512_DS && version != VERSION_SDXS_09; + } + void configure_weight_loading() { + const auto* sd_ctx_params = &config_->params; + const auto& model_loader = model_manager->loader(); + const auto wtype_stat = model_loader.get_wtype_stat(); bool have_int8_tensorwise = false; for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { if (tensor_storage.is_int8_tensorwise) { @@ -1022,638 +1224,122 @@ class StableDiffusionGGML { if (enable_mmap && apply_lora_immediately) { LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap"); } - model_loader.process_model_files(enable_mmap, needs_writable_mmap); - load_alphas_cumprod(model_loader); - - diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; + model_manager->prepare_file_io(); + load_alphas_cumprod(); + } - size_t text_encoder_params_mem_size = 0; - size_t unet_params_mem_size = 0; - size_t vae_params_mem_size = 0; - control_net_params_mem_size = 0; - size_t extension_params_mem_size = 0; + sd::model_builders::Context model_build_context() { + return {config_->params, version, model_manager->loader().get_tensor_storage_map(), backend_manager, model_manager}; + } - bool tae_preview_only = sd_ctx_params->tae_preview_only; - if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) { - tae_preview_only = false; - use_tae = true; + bool build_core_runners() { + sd::model_builders::CoreRunners runners; + if (!sd::model_builders::build_core_runners(model_build_context(), runners)) { + return false; } - // if (sd_version_is_minimax_h3(version) && use_tae) { - // LOG_WARN("MiniMax-H3 does not have a compatible TAE; ignoring --taesd"); - // tae_preview_only = false; - // use_tae = false; - // } - - auto& tensor_storage_map = model_loader.get_tensor_storage_map(); - - { - if (!ensure_backend_pair(SDBackendModule::TE) || - !ensure_backend_pair(SDBackendModule::DIFFUSION)) { - return false; - } - - if (sd_version_is_sd3(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_pid(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.net", - model_manager); - } else if (sd_version_is_ideogram4(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_krea2(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_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) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - false, - 1, - false, - model_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_OVIS_IMAGE) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - } else { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - model_manager); - } - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { - bool is_chroma = false; - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_ltxav(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - "text_encoders.llm", - "text_embedding_projection", - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_minimax_h3(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_hunyuan_video(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager); - } else if (sd_version_is_wan(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - true, - 0, - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager); - if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { - high_noise_diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.high_noise_diffusion_model", - version, - model_manager); - } - if (diffusion_model->get_desc() == "Wan2.1-I2V-14B" || - diffusion_model->get_desc() == "Wan2.1-FLF2V-14B" || - diffusion_model->get_desc() == "Wan2.1-I2V-1.3B") { - if (!ensure_backend_pair(SDBackendModule::CLIP_VISION)) { - return false; - } - clip_vision = std::make_shared(backend_for(SDBackendModule::CLIP_VISION), - tensor_storage_map, - model_manager); - clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); - if (!register_runner_params("CLIP vision", - clip_vision, - SDBackendModule::CLIP_VISION)) { - return false; - } - } - } 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; - } - } - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_qwen_image(version)) { - bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_mage_flow(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_longcat(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_HIDREAM_O1) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model", - model_manager); - } else if (sd_version_is_minit2i(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.model.net", - model_manager); - } else if (sd_version_is_anima(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_z_image(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager); - } else if (sd_version_is_boogu_image(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager); - } else if (sd_version_is_ernie_image(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else if (sd_version_is_lens(version)) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - model_manager); - } else { // SD1.x SD2.x SDXL - std::map 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)); - } - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - tensor_storage_map, - embbeding_map, - version, - model_manager); - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - model_manager); - if (sd_ctx_params->diffusion_conv_direct) { - LOG_INFO("Using Conv2d direct in the diffusion model"); - diffusion_model->set_conv2d_direct_enabled(true); - } - } - - cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE)); - if (!register_runner_params("Conditioner model", - cond_stage_model, - SDBackendModule::TE, - &text_encoder_params_mem_size)) { - return false; - } - - diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - if (!register_runner_params("Diffusion model", - diffusion_model, - SDBackendModule::DIFFUSION, - &unet_params_mem_size)) { - return false; - } + cond_stage_model = std::move(runners.conditioner); + diffusion_model = std::move(runners.diffusion); + high_noise_diffusion_model = std::move(runners.high_noise_diffusion); + clip_vision = std::move(runners.clip_vision); + ip_adapter = std::move(runners.ip_adapter); - if (high_noise_diffusion_model) { - high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - if (!register_runner_params("High noise diffusion model", - high_noise_diffusion_model, - SDBackendModule::DIFFUSION, - &unet_params_mem_size)) { - return false; - } - } + cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE)); + diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); + } + if (clip_vision) { + clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); + } + return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) && + register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) && + register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) && + register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) && + register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION); + } + + bool build_vae_runners() { + sd::model_builders::VAEOptions options; + options.use_tae = uses_tae(); + options.tae_preview_only = tae_preview_only(); + options.use_audio_vae = config_->use_audio_vae; + options.external_vae_is_invalid = external_vae_is_invalid; + sd::model_builders::VAERunners runners; + if (!sd::model_builders::build_vae_runners(model_build_context(), options, runners)) { + return false; + } + first_stage_model = std::move(runners.vae); + preview_vae = std::move(runners.preview); + audio_vae_model = std::move(runners.audio); - if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && clip_vision == nullptr) { - if (!ensure_backend_pair(SDBackendModule::CLIP_VISION)) { - return false; - } - clip_vision = std::make_shared(backend_for(SDBackendModule::CLIP_VISION), - tensor_storage_map, - model_manager); - clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); - if (!register_runner_params("CLIP vision", - clip_vision, - SDBackendModule::CLIP_VISION)) { - return false; - } - } + first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); + if (preview_vae) { + preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); + } + return register_runner_params(ModelComponent::VAE, first_stage_model, SDBackendModule::VAE) && + register_runner_params(ModelComponent::PreviewVAE, preview_vae, SDBackendModule::VAE) && + register_runner_params(ModelComponent::AudioVAE, audio_vae_model, SDBackendModule::VAE); + } - if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { - ip_adapter = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - tensor_storage_map, - "ip_adapter", - model_manager); - if (!register_runner_params("IP-Adapter", - ip_adapter, - SDBackendModule::DIFFUSION)) { - return false; - } - } + bool build_control_net_runner() { + if (config_->control_net_file == 0) { + return true; + } + if (!sd::model_builders::build_control_net_runner(model_build_context(), control_net)) { + return false; + } + control_net->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); + return register_runner_params(ModelComponent::ControlNet, control_net, SDBackendModule::CONTROL_NET); + } - if (!ensure_backend_pair(SDBackendModule::VAE)) { + bool build_extension_runners() { + GenerationExtensionInitContext extension_ctx{ + &config_->params, + version, + model_manager->loader().get_tensor_storage_map(), + config_->photomaker_source_available, + model_manager, + n_threads, + [this](SDBackendModule module) { return ensure_backend_pair(module); }, + [this](SDBackendModule module) { return backend_for(module); }, + [this](SDBackendModule module) { return params_backend_for(module); }, + }; + if (!sd::model_builders::build_extension_runners(extension_ctx, generation_extensions)) { + return false; + } + for (auto& extension : generation_extensions) { + if (!register_runner_params(extension->component(), extension, SDBackendModule::PHOTOMAKER)) { return false; } + } + return true; + } - auto create_tae = [&](bool decode_only) -> std::shared_ptr { - 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(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "decoder", - decode_only, - version, - model_manager); - - } else { - auto model = std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "decoder.layers", - decode_only, - version, - model_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 { - if (sd_version_is_ltxav(version)) { - return std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - version, - model_manager); - } else if (sd_version_is_minimax_h3(version)) { - return std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - model_manager); - } else if (sd_version_is_mage_flow(vae_version)) { - return std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - model_manager); - } else if (sd_version_uses_hunyuan_video_vae(vae_version)) { - return std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - vae_version, - model_manager); - } else if (sd_version_uses_wan_vae(vae_version)) { - return std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - vae_version, - model_manager); - } else { - auto model = std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - false, - vae_version, - model_manager); - if (sd_version_is_sdxl(version) && - (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || 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"); - first_stage_model = std::make_shared(version, - backend_for(SDBackendModule::VAE), - model_manager); - if (!register_runner_params("VAE", - first_stage_model, - SDBackendModule::VAE, - &vae_params_mem_size)) { - return false; - } - } else if (use_tae && !tae_preview_only) { - LOG_INFO("using TAE for encoding / decoding"); - first_stage_model = create_tae(false); - first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); - if (!register_runner_params("VAE", - first_stage_model, - SDBackendModule::VAE, - &vae_params_mem_size)) { - return false; - } - } else { - LOG_INFO("using VAE for encoding / decoding"); - first_stage_model = create_vae(); - first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); - if (!register_runner_params("VAE", - first_stage_model, - SDBackendModule::VAE, - &vae_params_mem_size)) { - return false; - } - if (use_tae && tae_preview_only) { - LOG_INFO("using TAE for preview"); - preview_vae = create_tae(true); - preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); - if (!register_runner_params("preview VAE", - preview_vae, - SDBackendModule::VAE, - &vae_params_mem_size)) { - return false; - } - } - } - - if (use_audio_vae) { - if (sd_version_is_minimax_h3(version)) { - audio_vae_model = std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "", - model_manager); - } else { - audio_vae_model = std::make_shared(backend_for(SDBackendModule::VAE), - tensor_storage_map, - "", - model_manager); - } - if (!register_runner_params("audio VAE", - audio_vae_model, - SDBackendModule::VAE, - &vae_params_mem_size)) { - return false; - } - } - - if (sd_ctx_params->vae_conv_direct) { - LOG_INFO("Using Conv2d direct in the vae model"); - first_stage_model->set_conv2d_direct_enabled(true); - if (preview_vae) { - preview_vae->set_conv2d_direct_enabled(true); - } - } - - if (use_control_net) { - if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) { - return false; - } - control_net = std::make_shared(backend_for(SDBackendModule::CONTROL_NET), - params_backend_for(SDBackendModule::CONTROL_NET), - model_loader.get_tensor_storage_map(), - version, - "", - model_manager); - control_net->set_max_graph_vram_bytes( - max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); - if (sd_ctx_params->diffusion_conv_direct) { - LOG_INFO("Using Conv2d direct in the control net"); - control_net->set_conv2d_direct_enabled(true); - } - if (!register_runner_params("ControlNet", - control_net, - SDBackendModule::CONTROL_NET, - &control_net_params_mem_size)) { - return false; - } - } - - { - generation_extensions.clear(); - auto photomaker_extension = create_photomaker_extension(); - GenerationExtensionInitContext extension_ctx{ - sd_ctx_params, - version, - tensor_storage_map, - model_loader, - model_manager, - n_threads, - [this](SDBackendModule module) { return ensure_backend_pair(module); }, - [this](SDBackendModule module) { return backend_for(module); }, - [this](SDBackendModule module) { return params_backend_for(module); }, - }; - if (!photomaker_extension->init(extension_ctx)) { - return false; - } - if (photomaker_extension->is_enabled()) { - generation_extensions.push_back(photomaker_extension); - } - - auto pulid_extension = create_pulid_extension(); - if (!pulid_extension->init(extension_ctx)) { - return false; - } - if (pulid_extension->is_enabled()) { - generation_extensions.push_back(pulid_extension); - } + bool validate_and_load_runners() { + const auto* sd_ctx_params = &config_->params; + const bool use_tae = uses_tae(); + const bool tae_preview_only = this->tae_preview_only(); + if (sd_ctx_params->flash_attn) { + LOG_INFO("Using flash attention"); + cond_stage_model->set_flash_attention_enabled(true); + if (clip_vision) { + clip_vision->set_flash_attention_enabled(true); } - for (auto& extension : generation_extensions) { - if (!register_runner_params(extension->name(), - extension, - SDBackendModule::PHOTOMAKER, - &extension_params_mem_size)) { - return false; - } + if (first_stage_model) { + first_stage_model->set_flash_attention_enabled(true); } - - if (sd_ctx_params->flash_attn) { - LOG_INFO("Using flash attention"); - cond_stage_model->set_flash_attention_enabled(true); - if (clip_vision) { - clip_vision->set_flash_attention_enabled(true); - } - if (first_stage_model) { - first_stage_model->set_flash_attention_enabled(true); - } - if (preview_vae) { - preview_vae->set_flash_attention_enabled(true); - } + if (preview_vae) { + preview_vae->set_flash_attention_enabled(true); } + } - if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) { - LOG_INFO("Using flash attention in the diffusion model"); - diffusion_model->set_flash_attention_enabled(true); - if (high_noise_diffusion_model) { - high_noise_diffusion_model->set_flash_attention_enabled(true); - } + if (sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) { + LOG_INFO("Using flash attention in the diffusion model"); + diffusion_model->set_flash_attention_enabled(true); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_flash_attention_enabled(true); } } - LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; @@ -1717,9 +1403,14 @@ class StableDiffusionGGML { } { - size_t total_params_ram_size = 0; - size_t total_params_vram_size = 0; - auto add_params_memory = [&](size_t size, SDBackendModule module) { + size_t text_encoder_params_mem_size = model_manager->registered_params_size({ModelComponent::Conditioner}); + size_t unet_params_mem_size = model_manager->registered_params_size({ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion}); + size_t vae_params_mem_size = model_manager->registered_params_size(runner_components().at(RunnerGroup::VAE)); + size_t control_net_params_mem_size = model_manager->registered_params_size({ModelComponent::ControlNet}); + size_t extension_params_mem_size = model_manager->registered_params_size(runner_components().at(RunnerGroup::Extensions)); + size_t total_params_ram_size = 0; + size_t total_params_vram_size = 0; + auto add_params_memory = [&](size_t size, SDBackendModule module) { if (size == 0) { return true; } @@ -1771,141 +1462,200 @@ class StableDiffusionGGML { extension_params_mem_size / 1024.0 / 1024.0, params_memory_location(extension_params_mem_size, SDBackendModule::PHOTOMAKER)); } + return true; + } - // init denoiser - { - prediction_t pred_type = sd_ctx_params->prediction; - - if (pred_type == PREDICTION_COUNT) { - if (sd_version_is_sd2(version)) { - pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; - } else if (sd_version_is_sdxl(version)) { - if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { - // CosXL models - // TODO: get sigma_min and sigma_max values from file - pred_type = EDM_V_PRED; - } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { - pred_type = V_PRED; - } else { - pred_type = EPS_PRED; - } - } else if (sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_hunyuan_video(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_qwen_image(version) || - sd_version_is_mage_flow(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_ernie_image(version) || - sd_version_is_z_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version)) { - pred_type = FLOW_PRED; - if (sd_version_is_wan(version)) { - default_flow_shift = 5.f; - } else if (sd_version_is_hunyuan_video(version)) { - default_flow_shift = 7.f; - } else if (sd_version_is_minimax_h3(version)) { - default_flow_shift = 12.f; - } else if (sd_version_is_ernie_image(version)) { - default_flow_shift = 4.f; - } else if (sd_version_is_pid(version)) { - default_flow_shift = 1.5f; - } else if (sd_version_is_ideogram4(version)) { - default_flow_shift = 1.0f; - } else if (sd_version_is_boogu_image(version)) { - default_flow_shift = 3.16f; - } else if (sd_version_is_mage_flow(version)) { - default_flow_shift = 6.f; - } else { - default_flow_shift = 3.f; - } - } else if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_longcat(version) || - sd_version_is_lens(version) || - sd_version_is_ltxav(version) || - sd_version_is_krea2(version)) { - pred_type = FLUX_FLOW_PRED; - - default_flow_shift = 1.0f; // TODO: validate - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { - default_flow_shift = 1.15f; - break; - } - } - if (sd_version_is_longcat(version)) { - default_flow_shift = 3.0f; - } else if (sd_version_is_lens(version)) { - default_flow_shift = 1.83f; - } else if (sd_version_is_ltxav(version)) { - default_flow_shift = 2.37f; - } else if (sd_version_is_krea2(version)) { + bool build_denoiser() { + const auto* sd_ctx_params = &config_->params; + const auto& model_loader = model_manager->loader(); + const auto& tensor_storage_map = model_loader.get_tensor_storage_map(); + denoiser = std::make_shared(); + default_flow_shift = INFINITY; + prediction_t pred_type = sd_ctx_params->prediction; + + if (pred_type == PREDICTION_COUNT) { + if (sd_version_is_sd2(version)) { + pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; + } else if (sd_version_is_sdxl(version)) { + if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { + // CosXL models + // TODO: get sigma_min and sigma_max values from file + pred_type = EDM_V_PRED; + } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { + pred_type = V_PRED; + } else { + pred_type = EPS_PRED; + } + } else if (sd_version_is_sd3(version) || + sd_version_is_wan(version) || + sd_version_is_hunyuan_video(version) || + sd_version_is_lingbot_video(version) || + sd_version_is_minimax_h3(version) || + sd_version_is_qwen_image(version) || + sd_version_is_mage_flow(version) || + version == VERSION_HIDREAM_O1 || + sd_version_is_anima(version) || + sd_version_is_ernie_image(version) || + sd_version_is_z_image(version) || + sd_version_is_boogu_image(version) || + sd_version_is_pid(version) || + sd_version_is_ideogram4(version)) { + pred_type = FLOW_PRED; + if (sd_version_is_wan(version)) { + default_flow_shift = 5.f; + } else if (sd_version_is_hunyuan_video(version)) { + default_flow_shift = 7.f; + } else if (sd_version_is_minimax_h3(version)) { + default_flow_shift = 12.f; + } else if (sd_version_is_ernie_image(version)) { + default_flow_shift = 4.f; + } else if (sd_version_is_pid(version)) { + default_flow_shift = 1.5f; + } else if (sd_version_is_ideogram4(version)) { + default_flow_shift = 1.0f; + } else if (sd_version_is_boogu_image(version)) { + default_flow_shift = 3.16f; + } else if (sd_version_is_mage_flow(version)) { + default_flow_shift = 6.f; + } else { + default_flow_shift = 3.f; + } + } else if (sd_version_is_flux(version) || + sd_version_is_flux2(version) || + sd_version_is_longcat(version) || + sd_version_is_lens(version) || + sd_version_is_ltxav(version) || + sd_version_is_krea2(version)) { + pred_type = FLUX_FLOW_PRED; + + default_flow_shift = 1.0f; // TODO: validate + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { default_flow_shift = 1.15f; + break; } - } else if (sd_version_is_sefi_image(version)) { - pred_type = SEFI_FLOW_PRED; - } else if (sd_version_is_minit2i(version)) { - pred_type = MINIT2I_FLOW_PRED; + } + if (sd_version_is_longcat(version)) { + default_flow_shift = 3.0f; + } else if (sd_version_is_lens(version)) { + default_flow_shift = 1.83f; + } else if (sd_version_is_ltxav(version)) { + default_flow_shift = 2.37f; + } else if (sd_version_is_krea2(version)) { + default_flow_shift = 1.15f; + } + } else if (sd_version_is_sefi_image(version)) { + pred_type = SEFI_FLOW_PRED; + } else if (sd_version_is_minit2i(version)) { + pred_type = MINIT2I_FLOW_PRED; + } else { + pred_type = EPS_PRED; + } + } + + switch (pred_type) { + case EPS_PRED: + LOG_INFO("running in eps-prediction mode"); + break; + case V_PRED: + LOG_INFO("running in v-prediction mode"); + denoiser = std::make_shared(); + break; + case EDM_V_PRED: + LOG_INFO("running in v-prediction EDM mode"); + denoiser = std::make_shared(); + break; + case FLOW_PRED: { + if (sd_version_is_ltxav(version)) { + LOG_INFO("running in LTXAV FLOW mode"); + denoiser = std::make_shared(); + } else if (sd_version_is_minimax_h3(version)) { + LOG_INFO("running in MiniMax H3 AV FLOW mode"); + denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); } else { - pred_type = EPS_PRED; + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); } + break; + } + case FLUX_FLOW_PRED: { + LOG_INFO("running in Flux FLOW mode"); + denoiser = std::make_shared(); + break; + } + case SEFI_FLOW_PRED: { + LOG_INFO("running in SeFi-Image dual-time FLOW mode"); + denoiser = std::make_shared(); + break; } + case MINIT2I_FLOW_PRED: { + LOG_INFO("running in MiniT2I FLOW mode"); + denoiser = std::make_shared(); + break; + } + default: { + LOG_ERROR("Unknown predition type %i", pred_type); + return false; + } + } - switch (pred_type) { - case EPS_PRED: - LOG_INFO("running in eps-prediction mode"); - break; - case V_PRED: - LOG_INFO("running in v-prediction mode"); - denoiser = std::make_shared(); - break; - case EDM_V_PRED: - LOG_INFO("running in v-prediction EDM mode"); - denoiser = std::make_shared(); - break; - case FLOW_PRED: { - if (sd_version_is_ltxav(version)) { - LOG_INFO("running in LTXAV FLOW mode"); - denoiser = std::make_shared(); - } else if (sd_version_is_minimax_h3(version)) { - LOG_INFO("running in MiniMax H3 AV FLOW mode"); - denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); - } else { - LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); - } + refresh_compvis_denoiser_sigmas(); + return true; + } + + bool build_runners(const RunnerGroups& groups) { + const auto& model_loader = model_manager->loader(); + std::map wtype_stat = model_loader.get_wtype_stat(); + std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); + std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); + std::map vae_wtype_stat = model_loader.get_vae_wtype_stat(); + + auto wtype_stat_to_str = [](const std::map& m, int key_width = 8, int value_width = 5) -> std::string { + std::ostringstream oss; + bool first = true; + for (const auto& [type, count] : m) { + if (!first) + oss << "|"; + first = false; + oss << std::right << std::setw(key_width) << ggml_type_name(type) + << ": " + << std::left << std::setw(value_width) << count; + } + return oss.str(); + }; + + LOG_INFO("Weight type stat: %s", wtype_stat_to_str(wtype_stat).c_str()); + LOG_INFO("Conditioner weight type stat: %s", wtype_stat_to_str(conditioner_wtype_stat).c_str()); + LOG_INFO("Diffusion model weight type stat: %s", wtype_stat_to_str(diffusion_model_wtype_stat).c_str()); + LOG_INFO("VAE weight type stat: %s", wtype_stat_to_str(vae_wtype_stat).c_str()); + + LOG_VERBOSE("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor)); + + configure_weight_loading(); + for (auto group : groups) { + bool success = false; + switch (group) { + case RunnerGroup::Core: + success = build_core_runners(); break; - } - case FLUX_FLOW_PRED: { - LOG_INFO("running in Flux FLOW mode"); - denoiser = std::make_shared(); + case RunnerGroup::VAE: + success = build_vae_runners(); break; - } - case SEFI_FLOW_PRED: { - LOG_INFO("running in SeFi-Image dual-time FLOW mode"); - denoiser = std::make_shared(); + case RunnerGroup::ControlNet: + success = build_control_net_runner(); break; - } - case MINIT2I_FLOW_PRED: { - LOG_INFO("running in MiniT2I FLOW mode"); - denoiser = std::make_shared(); + case RunnerGroup::Extensions: + success = build_extension_runners(); break; - } - default: { - LOG_ERROR("Unknown predition type %i", pred_type); - return false; - } } - - refresh_compvis_denoiser_sigmas(); + if (!success) { + return false; + } } - - return true; + if (!validate_and_load_runners()) { + return false; + } + return groups.count(RunnerGroup::Core) == 0 || build_denoiser(); } bool is_using_v_parameterization_for_sd2(bool is_inpaint = false) { @@ -1960,19 +1710,19 @@ class StableDiffusionGGML { if (lora_spec.is_high_noise) { LOG_VERBOSE("high noise lora: %s", lora_spec.path.c_str()); } - auto lora = std::make_shared(lora_log_id(lora_spec), - backend_for(module), - backend_for(module), - lora_spec.path, - lora_spec.is_high_noise ? "model.high_noise_" : "", - version); + const auto mode = backend_manager.params_backend_is_disk(module) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend; + auto lora = std::make_shared(lora_log_id(lora_spec), backend_for(module), params_backend_for(module), + model_manager, lora_spec.file_id, version, mode, + backend_manager.params_backend_follows_runtime(module)); LoraModel::filter_t lora_tensor_filter = module_filter; if (!lora_spec.tensor_name_prefix_filter.empty()) { lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) { return starts_with(tensor_name, prefix) && (!module_filter || module_filter(tensor_name)); }; } - 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()); return nullptr; } @@ -1999,13 +1749,20 @@ class StableDiffusionGGML { std::vector> load_runtime_loras_for_module(const std::vector& loras, const std::set& model_tensor_names, SDBackendModule module, - LoraModel::filter_t module_filter = nullptr) { + LoraModel::filter_t module_filter, + bool& success, + std::vector& next_models) { std::vector> module_lora_models; for (const auto& lora_spec : loras) { - auto lora = load_lora_model(lora_spec, module, module_filter); + auto cached = std::find_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) { + return entry.model != nullptr && entry.module == module && entry.matches(lora_spec); + }); + auto lora = cached == runtime_lora_models.end() ? load_lora_model(lora_spec, module, module_filter) + : std::move(cached->model); if (lora == nullptr) { if (lora_spec.required) { LOG_ERROR("required lora load failed: %s", lora_spec.path.c_str()); + success = false; } continue; } @@ -2014,36 +1771,43 @@ class StableDiffusionGGML { } lora->preprocess_lora_tensors(model_tensor_names); - runtime_lora_models.push_back(lora); + lora->multiplier = lora_spec.multiplier; + next_models.push_back({lora_spec, module, lora}); module_lora_models.push_back(std::move(lora)); } return module_lora_models; } - void apply_loras_immediately(const std::vector& loras) { + bool apply_loras_immediately(const std::vector& loras) { if (model_manager == nullptr) { if (!loras.empty()) { LOG_WARN("model manager is not available for immediate lora"); } - return; + return false; } clear_lora_adapters(); runtime_lora_models.clear(); - model_manager->set_loras(loras, version); + if (!loras.empty()) { + LOG_INFO("apply lora immediately"); + } + return model_manager->set_loras(loras, version); } - void apply_loras_at_runtime(const std::vector& loras) { + bool apply_loras_at_runtime(const std::vector& loras) { if (model_manager != nullptr) { - model_manager->set_loras({}, version); + if (!model_manager->set_loras({}, version)) + return false; } - runtime_lora_models.clear(); clear_lora_adapters(); if (loras.empty()) { - return; + runtime_lora_models.clear(); + return true; } + bool success = true; + std::vector next_models; std::set model_tensor_names; if (model_manager != nullptr) { model_tensor_names = model_manager->tensor_names(); @@ -2061,7 +1825,7 @@ class StableDiffusionGGML { load_runtime_loras_for_module(loras, model_tensor_names, SDBackendModule::TE, - lora_tensor_filter); + lora_tensor_filter, success, next_models); // Only attach the adapter when there are LoRAs targeting the cond_stage model. // An empty MultiLoraAdapter still routes every linear/conv through // forward_with_lora() instead of the direct kernel path — slower for no benefit. @@ -2081,7 +1845,7 @@ class StableDiffusionGGML { load_runtime_loras_for_module(loras, model_tensor_names, SDBackendModule::DIFFUSION, - lora_tensor_filter); + lora_tensor_filter, success, next_models); if (!diffusion_lora_models.empty()) { auto multi_lora_adapter = std::make_shared(diffusion_lora_models); diffusion_model->set_weight_adapter(multi_lora_adapter); @@ -2102,24 +1866,26 @@ class StableDiffusionGGML { load_runtime_loras_for_module(loras, model_tensor_names, SDBackendModule::VAE, - lora_tensor_filter); + lora_tensor_filter, success, next_models); if (!first_stage_lora_models.empty()) { auto multi_lora_adapter = std::make_shared(first_stage_lora_models); first_stage_model->set_weight_adapter(multi_lora_adapter); } } + runtime_lora_models = std::move(next_models); + return success; } void lora_stat() { if (!runtime_lora_models.empty()) { LOG_INFO("runtime_lora_models:"); for (auto& lora_model : runtime_lora_models) { - lora_model->stat(); + lora_model.model->stat(); } } } - void apply_loras(const sd_lora_t* loras, uint32_t lora_count) { + bool apply_loras(const sd_lora_t* loras, uint32_t lora_count) { std::vector all_loras; all_loras.reserve(lora_count); for (uint32_t i = 0; i < lora_count; i++) { @@ -2140,15 +1906,29 @@ class StableDiffusionGGML { } int64_t t0 = ggml_time_ms(); - if (apply_lora_immediately) { - apply_loras_immediately(all_loras); - } else { - apply_loras_at_runtime(all_loras); + end_runners(); + clear_lora_adapters(); + if (!model_manager->prepare_lora_sources(all_loras)) + return false; + runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) { + return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) { + return entry.matches(spec); + }); + }), + runtime_lora_models.end()); + const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras) + : apply_loras_at_runtime(all_loras); + if (!success) { + clear_lora_adapters(); + runtime_lora_models.clear(); + return false; } - int64_t t1 = ggml_time_ms(); + runner_state_.catalog_revision = model_manager->loader().revision(); + int64_t t1 = ggml_time_ms(); if (!all_loras.empty()) { LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); } + return true; } void reset_generation_extensions() { @@ -2748,7 +2528,7 @@ class StableDiffusionGGML { if (sd_version_is_unet(version)) { int nvf = -1; - if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) { + if (config_->animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) { nvf = static_cast(noised_input.shape()[3]); } UNetDiffusionExtra unet_extra{nvf, &controls, control_strength}; @@ -3505,22 +3285,6 @@ const char* sd_vae_format_name(enum sd_vae_format_t format) { } } -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; - } -} - void sd_cache_params_init(sd_cache_params_t* cache_params) { *cache_params = {}; cache_params->mode = SD_CACHE_DISABLED; @@ -3957,7 +3721,7 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { return false; } - if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { + if (sd_ctx->sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { return true; } return sd_version_supports_video_generation(sd_ctx->sd->version); @@ -5620,10 +5384,30 @@ static std::vector make_hires_sigma_schedule(sd_ctx_t* sd_ctx, sigmas.end()); } +static bool generate_image_impl(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* params, + sd_image_t** images_out, + int* num_images_out); + SD_API bool generate_image(sd_ctx_t* sd_ctx, - const sd_img_gen_params_t* sd_img_gen_params, + const sd_img_gen_params_t* params, sd_image_t** images_out, int* num_images_out) { + if (images_out != nullptr) + *images_out = nullptr; + if (num_images_out != nullptr) + *num_images_out = 0; + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || params == nullptr) { + return false; + } + StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); + return execution.ready && generate_image_impl(sd_ctx, params, images_out, num_images_out); +} + +static bool generate_image_impl(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out) { if (images_out != nullptr) { *images_out = nullptr; } @@ -5656,7 +5440,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, sd_ctx->sd->rng->manual_seed(request.seed); sd_ctx->sd->sampler_rng->manual_seed(request.seed); sd_ctx->sd->set_flow_shift(sd_img_gen_params->sample_params.flow_shift); - sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count); + if (!sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count)) + return false; apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y); const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args); @@ -6642,7 +6427,7 @@ static sd::Tensor upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx, const char* model_path, const sd::Tensor& packed_latent, int audio_length) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || packed_latent.empty()) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->model_manager == nullptr || packed_latent.empty()) { return {}; } if (strlen(SAFE_STR(model_path)) == 0) { @@ -6668,24 +6453,50 @@ static sd::Tensor upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx, (int)video_latent.shape()[3]); sd::Tensor unnormalized = sd_ctx->sd->un_normalize_ltx_video_latents(video_latent); + if (sd_ctx->sd->first_stage_model) { + sd_ctx->sd->first_stage_model->runner_end(); + } if (unnormalized.empty()) { LOG_ERROR("LTX latent un-normalization failed before spatial upscale"); return {}; } - auto upsampler_manager = std::make_shared(); - upsampler_manager->set_n_threads(sd_ctx->sd->n_threads); - upsampler_manager->set_enable_mmap(sd_ctx->sd->enable_mmap); - ModelLoader& model_loader = upsampler_manager->loader(); - if (!model_loader.init_from_file(model_path)) { + auto model_manager = sd_ctx->sd->model_manager; + struct UpsamplerScope { + ModelManager& manager; + ModelLoader::FileId owned_source = 0; + std::unique_ptr runner; + std::vector params; + + ~UpsamplerScope() { + if (runner) { + runner->runner_end(); + } + GGML_ASSERT(manager.unregister_param_tensors(params)); + if (owned_source != 0) { + GGML_ASSERT(manager.del_file(owned_source)); + } + } + } scope{*model_manager}; + + const std::string prefix = "ltx_latent_upsampler"; + ModelLoader candidate = model_manager->loader(); + ModelLoader::FileId source_file = 0; + if (!candidate.add_file(model_path, prefix + ".", &source_file)) { LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", model_path); return {}; } + const bool owns_source = model_manager->loader().file_revision(source_file) == 0; + if (!model_manager->set_loader(std::move(candidate))) { + return {}; + } + scope.owned_source = owns_source ? source_file : 0; - std::unique_ptr upsampler = - std::make_unique(sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), - model_loader.get_tensor_storage_map(), - upsampler_manager); + auto& upsampler = scope.runner; + upsampler = std::make_unique(sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), + model_manager->loader().get_tensor_storage_map(), + prefix, + model_manager); const size_t max_graph_vram_bytes = sd_ctx->sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (upsampler->model == nullptr) { @@ -6695,25 +6506,28 @@ static sd::Tensor upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx, std::map tensors; upsampler->get_param_tensors(tensors); - if (!upsampler_manager->register_param_tensors("LTX latent upsampler", - std::move(tensors), - ModelManager::ResidencyMode::ParamBackend, - sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), - sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)) || - !upsampler_manager->validate_registered_tensors()) { + for (const auto& entry : tensors) { + scope.params.push_back(entry.second); + } + if (!model_manager->register_param_tensors(ModelComponent::LatentUpsampler, + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), + sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)) || + !model_manager->validate_registered_tensors()) { LOG_ERROR("register LTX latent upsampler tensors with model manager failed"); return {}; } sd::Tensor upscaled = upsampler->compute(sd_ctx->sd->n_threads, unnormalized); - upsampler_manager.reset(); - upsampler.reset(); + upsampler->runner_end(); if (upscaled.empty()) { LOG_ERROR("LTX latent spatial upscale failed"); return {}; } upscaled = sd_ctx->sd->normalize_ltx_video_latents(upscaled); + sd_ctx->sd->first_stage_model->runner_end(); if (upscaled.empty()) { LOG_ERROR("LTX latent normalization failed after spatial upscale"); return {}; @@ -6862,7 +6676,7 @@ static bool generate_animatediff_video(sd_ctx_t* sd_ctx, img_gen_params.circular_y = sd_vid_gen_params->circular_y; sd_ctx->sd->animatediff_num_frames = n_frames; - bool ok = generate_image(sd_ctx, &img_gen_params, frames_out, num_frames_out); + bool ok = generate_image_impl(sd_ctx, &img_gen_params, frames_out, num_frames_out); sd_ctx->sd->animatediff_num_frames = 0; return ok; } @@ -6872,9 +6686,10 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_image_t** frames_out, int* num_frames_out, sd_audio_t** audio_out) { - if (sd_ctx == nullptr || sd_vid_gen_params == nullptr) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { return false; } + if (frames_out != nullptr) { *frames_out = nullptr; } @@ -6885,7 +6700,12 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, *num_frames_out = 0; } - if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { + StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); + if (!execution.ready) { + return false; + } + + if (sd_ctx->sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); return generate_animatediff_video(sd_ctx, sd_vid_gen_params, frames_out, num_frames_out); @@ -6919,7 +6739,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_ctx->sd->rng->manual_seed(request.seed); sd_ctx->sd->sampler_rng->manual_seed(request.seed); sd_ctx->sd->set_flow_shift(sd_vid_gen_params->sample_params.flow_shift); - sd_ctx->sd->apply_loras(sd_vid_gen_params->loras, sd_vid_gen_params->lora_count); + if (!sd_ctx->sd->apply_loras(sd_vid_gen_params->loras, sd_vid_gen_params->lora_count)) + return false; sd_ctx->sd->reset_generation_extensions(); SamplePlan plan(sd_ctx, sd_vid_gen_params, request); diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 46d4e4f3a..6b7b0efb0 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -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 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),