Compare commits

...

4 Commits

50 changed files with 8589 additions and 7048 deletions

View File

@ -229,6 +229,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
"src/model/*/*.h" "src/model/*/*.h"
"src/model/*/*.cpp" "src/model/*/*.cpp"
"src/model/*/*.hpp" "src/model/*/*.hpp"
"src/pipeline/*.h"
"src/pipeline/*.cpp"
"src/runtime/*.h" "src/runtime/*.h"
"src/runtime/*.cpp" "src/runtime/*.cpp"
"src/runtime/*.hpp" "src/runtime/*.hpp"

View File

@ -188,8 +188,9 @@ weights, compute buffers and caches must
still fit the runner's capacity checks. Offloading weights does not guarantee still fit the runner's capacity checks. Offloading weights does not guarantee
that every resolution or frame count will fit, and auto-fit does not change a that every resolution or frame count will fit, and auto-fit does not change a
component to CPU computation solely because its full weights exceed VRAM. component to CPU computation solely because its full weights exceed VRAM.
If a VAE decode fails, auto-fit retries with spatial tiling; supported video If a VAE decode fails, decoding retries with spatial tiling even when `--auto-fit`
decoders try temporal tiling first and can then add spatial tiling. is off; supported video decoders try temporal tiling first and can then add
spatial tiling. Spatial retries use half-size tiles along each latent dimension.
## Modules ## Modules

View File

@ -57,7 +57,7 @@ The RPC server acts as the worker. You must explicitly enable the **backend** (t
To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository. To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository.
> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=128`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled. > **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=160`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
> >
> I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts. > I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts.
@ -72,8 +72,8 @@ cmake .. -DGGML_RPC=ON \
-DGGML_VULKAN=ON \ # Ensure backend is enabled -DGGML_VULKAN=ON \ # Ensure backend is enabled
-DGGML_BUILD_SHARED_LIBS=OFF \ -DGGML_BUILD_SHARED_LIBS=OFF \
-DLLAMA_CURL=OFF \ -DLLAMA_CURL=OFF \
-DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \ -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
-DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
cmake --build . --config Release --target rpc-server -j $(nproc) cmake --build . --config Release --target rpc-server -j $(nproc)
``` ```
@ -86,8 +86,8 @@ cmake .. -DGGML_RPC=ON \
-DGGML_METAL=ON \ -DGGML_METAL=ON \
-DGGML_BUILD_SHARED_LIBS=OFF \ -DGGML_BUILD_SHARED_LIBS=OFF \
-DLLAMA_CURL=OFF \ -DLLAMA_CURL=OFF \
-DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \ -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 \
-DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
cmake --build . --config Release --target rpc-server cmake --build . --config Release --target rpc-server
``` ```
@ -101,8 +101,8 @@ cmake .. -G "Visual Studio 17 2022" -A x64 `
-DGGML_VULKAN=ON ` -DGGML_VULKAN=ON `
-DGGML_BUILD_SHARED_LIBS=OFF ` -DGGML_BUILD_SHARED_LIBS=OFF `
-DLLAMA_CURL=OFF ` -DLLAMA_CURL=OFF `
-DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 ` -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=160 `
-DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=160
cmake --build . --config Release --target rpc-server cmake --build . --config Release --target rpc-server
``` ```

View File

@ -11,6 +11,8 @@ $patterns = @(
"src/extensions/*.cpp" "src/extensions/*.cpp"
"src/extensions/*.h" "src/extensions/*.h"
"src/extensions/*.hpp" "src/extensions/*.hpp"
"src/pipeline/*.cpp"
"src/pipeline/*.h"
"src/runtime/*.cpp" "src/runtime/*.cpp"
"src/runtime/*.h" "src/runtime/*.h"
"src/runtime/*.hpp" "src/runtime/*.hpp"

View File

@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \
src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \ src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
src/core/*.cpp src/core/*.h src/core/*.hpp \ src/core/*.cpp src/core/*.h src/core/*.hpp \
src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \ src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \
src/pipeline/*.cpp src/pipeline/*.h \
src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \ src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \
src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \ src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \
src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \ src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \

View File

@ -390,6 +390,8 @@ namespace sd::backend_fit {
retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal"; retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal";
} else if (!tiling_params.enabled) { } else if (!tiling_params.enabled) {
tiling_params.enabled = true; tiling_params.enabled = true;
tiling_params.rel_size_x = 0.5f;
tiling_params.rel_size_y = 0.5f;
if (tiling_params.tile_size_x <= 0) { if (tiling_params.tile_size_x <= 0) {
tiling_params.tile_size_x = 256; tiling_params.tile_size_x = 256;
} }
@ -401,7 +403,7 @@ namespace sd::backend_fit {
return false; return false;
} }
LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
retry_mode); retry_mode);
return true; return true;
} }

View File

@ -9,7 +9,7 @@
#define EPS 1e-05f #define EPS 1e-05f
static_assert(GGML_MAX_NAME >= 128, "GGML_MAX_NAME must be at least 128"); static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160");
// n-mode tensor-matrix product // n-mode tensor-matrix product
// example: 2-mode product // example: 2-mode product

View File

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

View File

@ -38,9 +38,19 @@ public:
insert(kv); insert(kv);
} }
OrderedMap(const OrderedMap&) = default; OrderedMap(const OrderedMap& other) {
OrderedMap(OrderedMap&&) noexcept = default; for (const auto& value : other) {
OrderedMap& operator=(const OrderedMap&) = default; 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; OrderedMap& operator=(OrderedMap&&) noexcept = default;
// --- element access --- // --- element access ---
@ -174,4 +184,4 @@ public:
} }
}; };
#endif // __SD_CORE_ORDERED_MAP_HPP__ #endif // __SD_CORE_ORDERED_MAP_HPP__

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -1064,13 +1064,14 @@ struct MMDiTRunner : public DiffusionModelRunner {
{ {
LOG_INFO("loading from '%s'", file_path.c_str()); 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)) { if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return; 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, *mmdit,
"model.diffusion_model", "model.diffusion_model",
ModelManager::ResidencyMode::ParamBackend, ModelManager::ResidencyMode::ParamBackend,

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -1494,13 +1494,14 @@ namespace WAN {
{ {
LOG_INFO("loading from '%s'", file_path.c_str()); 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.")) { 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()); LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return; 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, *vae,
ModelManager::ResidencyMode::ParamBackend, ModelManager::ResidencyMode::ParamBackend,
backend, backend,

61
src/model_component.h Normal file
View File

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

View File

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

View File

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

View File

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

View File

@ -9,6 +9,7 @@
#include <mutex> #include <mutex>
#include <regex> #include <regex>
#include <set> #include <set>
#include <stdexcept>
#include <string> #include <string>
#include <thread> #include <thread>
#include <unordered_map> #include <unordered_map>
@ -27,6 +28,7 @@
#include "ggml-alloc.h" #include "ggml-alloc.h"
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
#include "json.hpp"
#include "zip.h" #include "zip.h"
#include "name_conversion.h" #include "name_conversion.h"
@ -151,15 +153,19 @@ ModelLoader::ModelLoader()
} }
size_t ModelLoader::add_file_path(const std::string& file_path) { size_t ModelLoader::add_file_path(const std::string& file_path) {
if (model_files_processed) { auto it = std::find(file_paths_.begin(), file_paths_.end(), file_path);
file_data.clear(); if (it != file_paths_.end()) {
model_files_processed = false; return static_cast<size_t>(it - file_paths_.begin());
} }
invalidate_file_data();
file_paths_.push_back(file_path); file_paths_.push_back(file_path);
return file_paths_.size() - 1; return file_paths_.size() - 1;
} }
void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) { 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; 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) { 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)) { if (is_directory(file_path)) {
LOG_INFO("load %s using diffusers format", file_path.c_str()); LOG_INFO("load %s using diffusers format", file_path.c_str());
return init_from_diffusers_file(file_path, prefix); 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() { void ModelLoader::convert_tensors_name() {
SDVersion version = (version_ == VERSION_COUNT) ? get_sd_version() : version_; if (names_converted_) {
String2TensorStorage new_map; return;
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);
} }
names_converted_ = true;
tensor_storage_map.swap(new_map); rebuild_catalog();
} }
bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, const std::string& prefix, SDVersion version) { bool ModelLoader::init_from_file_and_convert_name(const std::string& file_path, const std::string& prefix, SDVersion version) {
@ -257,7 +266,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
std::vector<TensorStorage> tensor_storages; std::vector<TensorStorage> tensor_storages;
std::string error; 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()); LOG_ERROR("%s", error.c_str());
return false; 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) { for (const std::string& shard_path : shard_paths) {
if (!init_from_file(shard_path, prefix)) { if (!parse_file(shard_path, prefix)) {
return false;
}
}
std::ifstream index_file(file_path);
const auto index = nlohmann::json::parse(index_file);
for (const auto& entry : index.at("weight_map").items()) {
const auto expected = (std::filesystem::u8path(file_path).parent_path() /
std::filesystem::u8path(entry.value().get<std::string>()))
.lexically_normal();
bool found = false;
for (const auto& shard : parsed_tensor_names_) {
if (std::filesystem::u8path(shard.first).lexically_normal() == expected) {
found = shard.second.count(entry.key()) != 0;
break;
}
}
if (!found) {
LOG_ERROR("safetensors index tensor '%s' is missing from its declared shard", entry.key().c_str());
return false; 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_path = path_join(file_path, "text_encoder/model.safetensors");
std::string clip_g_path = path_join(file_path, "text_encoder_2/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; return false;
} }
if (!init_from_safetensors_file(vae_path, "vae.")) { if (file_exists(vae_path) && !parse_file(vae_path, prefix + "vae.")) {
LOG_WARN("Couldn't find working VAE in %s", file_path.c_str()); return false;
// return false;
} }
if (!init_from_safetensors_file(clip_path, "te.")) { if (file_exists(clip_path) && !parse_file(clip_path, prefix + "te.")) {
LOG_WARN("Couldn't find working text encoder in %s", file_path.c_str()); return false;
// return false;
} }
if (!init_from_safetensors_file(clip_g_path, "te.1.")) { if (file_exists(clip_g_path) && !parse_file(clip_g_path, prefix + "te.1.")) {
LOG_VERBOSE("Couldn't find working second text encoder in %s", file_path.c_str()); return false;
} }
return true; return true;
} }
SDVersion ModelLoader::get_sd_version() { SDVersion ModelLoader::get_sd_version() const {
TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight; TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
bool has_multiple_encoders = false; bool has_multiple_encoders = false;
@ -623,7 +649,7 @@ SDVersion ModelLoader::get_sd_version() {
return VERSION_COUNT; return VERSION_COUNT;
} }
std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() { std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat; std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) { for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) { if (is_unused_tensor(tensor_storage.name)) {
@ -640,7 +666,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_wtype_stat() {
return wtype_stat; return wtype_stat;
} }
std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() { std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat; std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) { for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) { if (is_unused_tensor(tensor_storage.name)) {
@ -664,7 +690,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_conditioner_wtype_stat() {
return wtype_stat; return wtype_stat;
} }
std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() { std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat; std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) { for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) { if (is_unused_tensor(tensor_storage.name)) {
@ -685,7 +711,7 @@ std::map<ggml_type, uint32_t> ModelLoader::get_diffusion_model_wtype_stat() {
return wtype_stat; return wtype_stat;
} }
std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() { std::map<ggml_type, uint32_t> ModelLoader::get_vae_wtype_stat() const {
std::map<ggml_type, uint32_t> wtype_stat; std::map<ggml_type, uint32_t> wtype_stat;
for (auto& [name, tensor_storage] : tensor_storage_map) { for (auto& [name, tensor_storage] : tensor_storage_map) {
if (is_unused_tensor(tensor_storage.name)) { 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) { 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) { 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) { for (const auto& tensor_type_rule : map_rules) {
std::regex pattern(tensor_type_rule.first); std::regex pattern(tensor_type_rule.first);
if (std::regex_search(name, pattern)) { 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; tensor_storage.expected_type = dst_type;
} }
invalidate_file_data();
++revision_;
} }
void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
@ -829,6 +860,13 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors, std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors, std::set<std::string> ignore_tensors,
bool writable_mmap) { bool writable_mmap) {
std::set<std::string> names;
for (const auto& entry : tensors) {
names.insert(entry.first);
}
if (!validate_sources(&names)) {
return {};
}
process_model_files(true, writable_mmap); process_model_files(true, writable_mmap);
std::vector<MmapTensorStore> result; std::vector<MmapTensorStore> result;
@ -919,6 +957,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool enable_mmap, bool enable_mmap,
const std::set<std::string>* target_tensor_names, const std::set<std::string>* target_tensor_names,
bool log_progress) { bool log_progress) {
if (!validate_sources(target_tensor_names)) {
return false;
}
process_model_files(enable_mmap, false); process_model_files(enable_mmap, false);
std::atomic<int64_t> read_time_ms(0); std::atomic<int64_t> read_time_ms(0);
@ -1242,7 +1283,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
(convert_time_ms.load() / (float)last_n_threads) / 1000.f, (convert_time_ms.load() / (float)last_n_threads) / 1000.f,
(copy_to_backend_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) { 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; 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.offset != tensor_storage.offset ||
current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) { 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()); LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
@ -1440,7 +1483,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
return true; 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; const std::string& name = tensor_storage.name;
if (tensor_storage.is_int8_tensorwise) { if (tensor_storage.is_int8_tensorwise) {
return false; return false;
@ -1478,7 +1521,7 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
return false; 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; size_t alignment = 128;
if (backend != nullptr) { if (backend != nullptr) {
alignment = ggml_backend_get_alignment(backend); alignment = ggml_backend_get_alignment(backend);

View File

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

336
src/model_loader_files.cpp Normal file
View File

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

View File

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

View File

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

154
src/model_manager_files.cpp Normal file
View File

@ -0,0 +1,154 @@
#include "model_manager.h"
#include <algorithm>
#include "core/util.h"
static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
return a.file_id == b.file_id && a.file_revision == b.file_revision &&
a.file_index == b.file_index && a.offset == b.offset && a.index_in_zip == b.index_in_zip &&
a.storage_key == b.storage_key && a.type == b.type && a.expected_type == b.expected_type &&
a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot &&
a.int8_convrot_group_size == b.int8_convrot_group_size;
}
void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
auto affected = states;
for (const auto& block : params_storage_blocks_) {
if (std::any_of(block->states.begin(), block->states.end(), [&](TensorState* state) { return states.count(state) != 0; })) {
affected.insert(block->states.begin(), block->states.end());
}
}
for (auto it = prefetch_blocks_.begin(); it != prefetch_blocks_.end();) {
if (std::any_of(it->second->states.begin(), it->second->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
free_prefetch_block(*it->second);
it = prefetch_blocks_.erase(it);
} else {
++it;
}
}
for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
if (std::any_of((*it)->staged_tensors.begin(), (*it)->staged_tensors.end(), [&](const auto& entry) { return affected.count(entry.first) != 0; })) {
ggml_backend_synchronize((*it)->compute_backend);
free_compute_staging_block(**it);
it = compute_staging_blocks_.erase(it);
} else {
++it;
}
}
for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
if (std::any_of((*it)->states.begin(), (*it)->states.end(), [&](TensorState* state) { return affected.count(state) != 0; })) {
free_params_storage_block(**it);
it = params_storage_blocks_.erase(it);
} else {
++it;
}
}
for (auto* state : affected) {
state->metadata_validated = false;
state->applied_lora_epoch = UINT64_MAX;
}
}
bool ModelManager::set_loader(ModelLoader loader) {
if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) {
return state->pin_count != 0;
})) {
LOG_ERROR("cannot update model sources during execution");
return false;
}
std::map<std::pair<ModelLoader::FileId, SDVersion>, String2TensorStorage> scoped;
auto sources_for = [&](const TensorState& state) -> const String2TensorStorage& {
if (state.source_file == 0)
return loader.get_tensor_storage_map();
auto key = std::make_pair(state.source_file, state.source_version);
auto found = scoped.find(key);
if (found == scoped.end())
found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first;
return found->second;
};
bool lora_changed = false;
for (const auto& spec : loras_) {
lora_changed |= loader.file_revision(spec.file_id) != spec.file_revision;
}
std::unordered_set<TensorState*> changed;
for (const auto& state : tensor_states_) {
const auto& sources = sources_for(*state);
auto source = sources.find(state->name);
const bool found = source != sources.end();
if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
(lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
changed.insert(state.get());
}
}
invalidate_sources(changed);
for (auto* state : changed) {
const auto& sources = sources_for(*state);
auto source = sources.find(state->name);
state->has_source = source != sources.end();
state->source = state->has_source ? source->second : TensorStorage{};
}
if (lora_changed) {
++current_lora_epoch_;
for (auto& spec : loras_)
spec.file_revision = loader.file_revision(spec.file_id);
}
model_loader_ = std::move(loader);
model_loader_.set_n_threads(n_threads_);
return true;
}
bool ModelManager::add_file(const std::string& path, const std::string& prefix, ModelLoader::FileId* id, bool force) {
ModelLoader candidate = model_loader_;
ModelLoader::FileId added_id;
if (!candidate.add_file(path, prefix, &added_id, force) || !set_loader(std::move(candidate))) {
return false;
}
if (id != nullptr) {
*id = added_id;
}
return true;
}
bool ModelManager::del_file(ModelLoader::FileId id) {
ModelLoader candidate = model_loader_;
return candidate.del_file(id) && set_loader(std::move(candidate));
}
bool ModelManager::refresh_files() {
ModelLoader candidate = model_loader_;
return candidate.refresh_files() && set_loader(std::move(candidate));
}
ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComponent>& components, const ModelLoader& loader) const {
ModelLoader::FileVersions versions;
const auto& sources = loader.get_tensor_storage_map();
for (const auto& state : tensor_states_) {
if (components.count(state->component) == 0) {
continue;
}
if (state->source_file != 0) {
versions[state->source_file] = loader.file_revision(state->source_file);
continue;
}
auto source = sources.find(state->name);
if (source != sources.end()) {
versions[source->second.file_id] = source->second.file_revision;
}
}
return versions;
}
size_t ModelManager::registered_params_size(const std::set<ModelComponent>& components) const {
size_t bytes = 0;
std::unordered_set<const ggml_tensor*> seen;
for (const auto& state : tensor_states_) {
if (components.count(state->component) != 0 && state->tensor != nullptr && seen.insert(state->tensor).second) {
bytes += ggml_nbytes(state->tensor);
}
}
return bytes;
}

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,482 @@
#ifndef __SD_PIPELINE_DIFFUSION_ENGINE_H__
#define __SD_PIPELINE_DIFFUSION_ENGINE_H__
#include <atomic>
#include <cmath>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <string>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "core/tensor.hpp"
#include "core/util.h"
#include "model/adapter/lora.hpp"
#include "model_builders.h"
#include "model_manager.h"
#include "stable-diffusion.h"
class RNG;
struct Denoiser;
struct LoraModel;
struct ConditionerParams;
struct SDCondition;
struct RefImageParams;
extern const char* model_version_to_str[];
static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
return version == VERSION_FLUX ||
sd_version_is_flux2(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_mage_flow(version) ||
sd_version_is_longcat(version) ||
sd_version_is_z_image(version) ||
sd_version_is_boogu_image(version);
}
class StableDiffusionGGML {
public:
SDBackendManager backend_manager;
SDVersion version;
bool external_vae_is_invalid = false;
bool circular_x = false;
bool circular_y = false;
std::shared_ptr<RNG> rng;
std::shared_ptr<RNG> sampler_rng = nullptr;
int n_threads = -1;
float default_flow_shift = INFINITY;
float active_flow_shift = INFINITY;
std::shared_ptr<Conditioner> cond_stage_model;
std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision; // for svd or wan2.1 i2v
std::shared_ptr<DiffusionModelRunner> diffusion_model;
std::shared_ptr<DiffusionModelRunner> high_noise_diffusion_model;
std::shared_ptr<VAE> first_stage_model;
std::shared_ptr<VAE> preview_vae;
std::shared_ptr<AudioVAERunner> audio_vae_model;
std::shared_ptr<ControlNet> control_net;
std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
sd::Tensor<float> ip_adapter_tokens;
sd::Tensor<float> ip_adapter_uncond_tokens;
float ip_adapter_strength = 1.0f;
std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
struct RuntimeLora {
ModelManager::LoraSpec spec;
SDBackendModule module;
std::shared_ptr<LoraModel> 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<RuntimeLora> runtime_lora_models;
bool apply_lora_immediately = false;
int animatediff_num_frames = 0;
std::string taesd_path;
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
bool enable_mmap = false;
sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
bool disable_prefetch = false;
bool disable_segmented_compute = false;
bool eager_load = false;
std::string backend_spec;
std::string params_backend_spec;
std::string split_mode_spec;
bool auto_fit_enabled = false;
bool diffusion_conv_direct = false;
bool is_using_v_parameterization = false;
bool is_using_edm_v_parameterization = false;
std::shared_ptr<ModelManager> model_manager;
enum class RunnerGroup { Core,
VAE,
ControlNet,
Extensions };
using RunnerGroups = std::set<RunnerGroup>;
struct ModelConfig {
sd_ctx_params_t params{};
std::list<std::string> strings;
std::vector<sd_embedding_t> 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<RunnerGroup, ModelLoader::FileVersions> sources;
};
std::recursive_mutex execution_mutex;
std::unique_ptr<ModelConfig> config_;
RunnerState runner_state_;
bool executing_ = false;
std::shared_ptr<Denoiser> denoiser;
std::vector<float> file_alphas_cumprod;
StableDiffusionGGML();
~StableDiffusionGGML();
static const std::map<RunnerGroup, std::set<ModelComponent>>& runner_components();
static RunnerGroups all_runner_groups();
ModelLoader::FileVersions runner_source_versions(RunnerGroup group, const ModelLoader& loader) const;
void capture_runner_sources();
void end_runners();
bool reset_runners(const RunnerGroups& groups);
bool refresh_model_sources();
bool apply_model_update(ModelLoader candidate,
std::unique_ptr<ModelConfig> next_config = nullptr,
RunnerGroups groups = {});
struct ContextOperation {
StableDiffusionGGML& sd;
std::unique_lock<std::recursive_mutex> 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 params_backend_for(SDBackendModule module);
std::atomic<sd_cancel_mode_t> cancellation_flag = SD_CANCEL_RESET;
void set_cancel_flag(enum sd_cancel_mode_t flag);
void reset_cancel_flag();
enum sd_cancel_mode_t get_cancel_flag();
size_t max_graph_vram_bytes_for_module(SDBackendModule module);
std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends);
bool ensure_backend_pair(SDBackendModule module);
template <typename T>
bool register_runner_params(ModelComponent component,
const std::shared_ptr<T>& model,
SDBackendModule module,
size_t* params_mem_size = nullptr);
template <typename T>
bool register_row_split_runner_params(ModelComponent component,
const std::shared_ptr<T>& model,
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size);
// Register graph-cut layer-split tensors on the primary backend first.
// The first real graph assigns each param tensor to a runtime backend
// before weights are loaded or staged.
template <typename T>
bool register_layer_split_runner_params(ModelComponent component,
const std::shared_ptr<T>& model,
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size);
bool unload_control_net();
bool load_control_net_from_file(const std::string& path);
void apply_circular_axes(bool circular_x, bool circular_y);
bool init_backend();
bool row_split_active();
bool graph_cut_layer_split_active();
std::shared_ptr<RNG> get_rng(rng_type_t rng_type);
void refresh_compvis_denoiser_sigmas();
void load_alphas_cumprod();
bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration);
bool init(const sd_ctx_params_t* sd_ctx_params);
bool uses_tae() const;
bool tae_preview_only() const;
void configure_weight_loading();
sd::model_builders::Context model_build_context();
bool build_core_runners();
bool build_vae_runners();
bool build_control_net_runner();
bool build_extension_runners();
bool validate_and_load_runners();
bool build_denoiser();
bool build_runners(const RunnerGroups& groups);
bool is_using_v_parameterization_for_sd2(bool is_inpaint = false);
static std::string lora_log_id(const ModelManager::LoraSpec& lora);
std::shared_ptr<LoraModel> load_lora_model(const ModelManager::LoraSpec& lora_spec,
SDBackendModule module,
LoraModel::filter_t module_filter = nullptr);
void clear_lora_adapters();
std::vector<std::shared_ptr<LoraModel>> load_runtime_loras_for_module(const std::vector<ModelManager::LoraSpec>& loras,
const std::set<std::string>& model_tensor_names,
SDBackendModule module,
LoraModel::filter_t module_filter,
bool& success,
std::vector<RuntimeLora>& next_models);
bool apply_loras_immediately(const std::vector<ModelManager::LoraSpec>& loras);
bool apply_loras_at_runtime(const std::vector<ModelManager::LoraSpec>& loras);
void lora_stat();
bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
void reset_generation_extensions();
void prepare_generation_extensions(const sd_pm_params_t& pm_params,
const sd_pulid_params_t& pulid_params,
ConditionerParams& condition_params,
int total_steps);
sd::Tensor<float> get_clip_vision_output(const sd::Tensor<float>& image,
bool return_pooled = true,
int clip_skip = -1,
bool zero_out_masked = false);
void compute_ip_adapter_tokens(const sd_image_t& image, float strength);
std::vector<float> process_timesteps(const std::vector<float>& timesteps,
const sd::Tensor<float>& init_latent,
const sd::Tensor<float>& denoise_mask,
int step);
std::vector<float> process_ltxav_video_timesteps(const std::vector<float>& timesteps,
const sd::Tensor<float>& init_latent,
const sd::Tensor<float>& denoise_mask);
void preview_image(int step,
const sd::Tensor<float>& latents,
enum SDVersion version,
preview_t preview_mode,
std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
void* step_callback_data,
bool is_noisy);
std::vector<float> prepare_sample_timesteps(float sigma,
int shifted_timestep);
void adjust_sample_step_scalings(int shifted_timestep,
const std::vector<float>& timesteps_vec,
float c_in,
float* c_skip,
float* c_out);
struct SamplePreviewContext {
sd_preview_cb_t callback = nullptr;
void* data = nullptr;
preview_t mode = PREVIEW_NONE;
};
SamplePreviewContext prepare_sample_preview_context();
void report_sample_progress(int step,
size_t total_steps,
bool terminal_sigma_is_zero,
int64_t* last_progress_us);
void compute_sample_controls(const sd::Tensor<float>& control_image,
const sd::Tensor<float>& noised_input,
const sd::Tensor<float>& timesteps_tensor,
const SDCondition& condition,
std::vector<sd::Tensor<float>>* controls);
sd::Tensor<float> sample(const std::shared_ptr<DiffusionModelRunner>& work_diffusion_model,
bool inverse_noise_scaling,
const sd::Tensor<float>& init_latent,
sd::Tensor<float> noise,
const SDCondition& cond,
const SDCondition& uncond,
const SDCondition& img_uncond,
const sd::Tensor<float>& control_image,
float control_strength,
const sd_guidance_params_t& guidance,
float eta,
int shifted_timestep,
sample_method_t method,
bool is_flow_denoiser,
const char* extra_sample_args,
const std::vector<float>& sigmas,
const std::vector<sd::Tensor<float>>& ref_latents,
const RefImageParams& ref_image_params,
const sd::Tensor<float>& denoise_mask,
const sd::Tensor<float>& vace_context,
float vace_strength,
int audio_length,
float frame_rate,
const sd_cache_params_t* cache_params,
bool preview_final_step,
const sd::Tensor<float>& video_positions = {});
int get_vae_scale_factor();
int get_diffusion_model_down_factor();
int get_latent_channel();
int get_image_channels() const;
int get_image_seq_len(int h, int w);
sd::Tensor<float> generate_init_latent(int width,
int height,
int frames = 1,
bool video = false);
int video_frames_to_latent_frames(int frames);
int latent_frames_to_video_frames(int latent_frames);
int align_video_frames(int frames);
sd::Tensor<float> encode_to_vae_latents(const sd::Tensor<float>& x);
sd::Tensor<float> encode_first_stage(const sd::Tensor<float>& x);
sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false);
sd::Tensor<float> normalize_ltx_video_latents(const sd::Tensor<float>& x);
sd::Tensor<float> un_normalize_ltx_video_latents(const sd::Tensor<float>& x);
sd::Tensor<float> decode_ltx_audio_latent(const sd::Tensor<float>& audio_latent);
void set_flow_shift(float flow_shift = INFINITY);
bool is_flow_denoiser();
std::string get_default_ref_image_preset(SDVersion version) const;
RefImageParams resolve_ref_image_params(const char* ref_image_args) const;
};
#endif // __SD_PIPELINE_DIFFUSION_ENGINE_H__

73
src/pipeline/generation.h Normal file
View File

@ -0,0 +1,73 @@
#ifndef __SD_PIPELINE_GENERATION_H__
#define __SD_PIPELINE_GENERATION_H__
#include "conditioning/conditioner.hpp"
#include "stable-diffusion.h"
class StableDiffusionGGML;
static inline bool sd_version_supports_animatediff(SDVersion version) {
return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX;
}
namespace sd::pipeline {
struct ImageGenerationLatents {
sd::Tensor<float> init_latent;
sd::Tensor<float> concat_latent;
sd::Tensor<float> img_uncond_concat_latent;
sd::Tensor<float> audio_latent;
sd::Tensor<float> video_positions;
sd::Tensor<float> control_image;
std::vector<sd::Tensor<float>> ref_images;
std::vector<sd::Tensor<float>> ref_latents;
std::vector<sd::Tensor<float>> reference_audio_latents;
std::vector<MiniMaxH3ReferenceBlock> minimax_reference_blocks;
std::vector<MiniMaxH3PresentationItem> minimax_presentation_refs;
std::vector<int32_t> keyframe_indices;
sd::Tensor<float> denoise_mask;
sd::Tensor<float> clip_vision_output;
sd::Tensor<float> vace_context;
int64_t ref_image_num = 0;
int64_t video_conditioning_frame_count = 0;
int64_t video_target_frame_count = 0;
int audio_length = 0;
};
struct ImageGenerationEmbeds {
SDCondition cond;
SDCondition uncond;
SDCondition img_uncond;
};
struct ConditionerRunnerEndOnExit {
Conditioner* conditioner = nullptr;
~ConditionerRunnerEndOnExit() {
if (conditioner != nullptr) {
conditioner->runner_end();
}
}
};
// Callers hold ExecutionScope; AnimateDiff reuses the image path within the same scope.
bool generate_image(StableDiffusionGGML* sd,
const sd_img_gen_params_t* sd_img_gen_params,
sd_image_t** images_out,
int* num_images_out);
bool generate_video(StableDiffusionGGML* sd,
const sd_vid_gen_params_t* sd_vid_gen_params,
sd_image_t** frames_out,
int* num_frames_out,
sd_audio_t** audio_out);
sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
const char* model_path,
const sd::Tensor<float>& packed_latent,
int audio_length);
sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels);
} // namespace sd::pipeline
#endif // __SD_PIPELINE_GENERATION_H__

1034
src/pipeline/image.cpp Normal file

File diff suppressed because it is too large Load Diff

View File

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

View File

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

471
src/pipeline/request.cpp Normal file
View File

@ -0,0 +1,471 @@
#include "request.h"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include "diffusion_engine.h"
#include "runtime/denoiser.hpp"
namespace sd::pipeline {
const char* sampling_methods_str[] = {
"Euler",
"Euler A",
"Heun",
"DPM2",
"DPM++ (2s)",
"DPM++ (2M)",
"modified DPM++ (2M)",
"iPNDM",
"iPNDM_v",
"LCM",
"DDIM \"trailing\"",
"TCD",
"Res Multistep",
"Res 2s",
"ER-SDE",
"Euler CFG++",
"Euler A CFG++",
"Euler GE",
"DPM++ (2M) SDE",
"DPM++ (2M) SDE BT",
"LMS",
};
static_assert(SAMPLE_METHOD_COUNT == sizeof(sampling_methods_str) / sizeof(sampling_methods_str[0]),
"\nnumber of elements in sampling_methods_str[] != SAMPLE_METHOD_COUNT");
static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {
return sd_version_is_inpaint_or_unet_edit(version) ||
(has_ref_images && sd_version_supports_ref_latent_img_cfg(version));
}
enum sample_method_t default_sample_method(const StableDiffusionGGML* sd) {
if (sd != nullptr) {
if (sd_version_is_pid(sd->version)) {
return LCM_SAMPLE_METHOD;
}
if (sd_version_is_dit(sd->version)) {
return EULER_SAMPLE_METHOD;
}
}
return EULER_A_SAMPLE_METHOD;
}
enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method) {
if (sd != nullptr) {
auto edm_v_denoiser = std::dynamic_pointer_cast<EDMVDenoiser>(sd->denoiser);
if (edm_v_denoiser) {
return EXPONENTIAL_SCHEDULER;
}
}
if (sample_method == LCM_SAMPLE_METHOD || sample_method == TCD_SAMPLE_METHOD) {
return LCM_SCHEDULER;
} else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) {
return SIMPLE_SCHEDULER;
} else if (sd != nullptr && sd_version_is_flux(sd->version)) {
return FLUX_SCHEDULER;
} else if (sd != nullptr && sd_version_is_flux2(sd->version)) {
return FLUX2_SCHEDULER;
} else if (sd != nullptr && sd_version_is_ltxav(sd->version)) {
return LTX2_SCHEDULER;
} else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) {
return LOGIT_NORMAL_SCHEDULER;
}
return DISCRETE_SCHEDULER;
}
static int64_t resolve_seed(int64_t seed) {
if (seed >= 0) {
return seed;
}
srand((int)time(nullptr));
return rand();
}
static enum sample_method_t resolve_sample_method(StableDiffusionGGML* sd, enum sample_method_t sample_method) {
if (sample_method == SAMPLE_METHOD_COUNT) {
return default_sample_method(sd);
}
return sample_method;
}
static scheduler_t resolve_scheduler(StableDiffusionGGML* sd,
scheduler_t scheduler,
enum sample_method_t sample_method) {
if (scheduler == SCHEDULER_COUNT) {
return default_scheduler(sd, sample_method);
}
return scheduler;
}
float resolve_eta(StableDiffusionGGML* sd,
float eta,
enum sample_method_t sample_method) {
if (eta == INFINITY) {
if (sd->version == VERSION_HIDREAM_O1) {
return 8.f;
}
switch (sample_method) {
case DDIM_TRAILING_SAMPLE_METHOD:
case TCD_SAMPLE_METHOD:
case RES_MULTISTEP_SAMPLE_METHOD:
case RES_2S_SAMPLE_METHOD:
return 0.0f;
case EULER_A_SAMPLE_METHOD:
case DPMPP2S_A_SAMPLE_METHOD:
case ER_SDE_SAMPLE_METHOD:
case EULER_A_CFG_PP_SAMPLE_METHOD:
case DPMPP2M_SDE_SAMPLE_METHOD:
case DPMPP2M_SDE_BT_SAMPLE_METHOD:
return 1.0f;
default:;
}
return 0.0f;
}
return eta;
}
GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params) {
prompt = SAFE_STR(sd_img_gen_params->prompt);
negative_prompt = SAFE_STR(sd_img_gen_params->negative_prompt);
width = sd_img_gen_params->width;
height = sd_img_gen_params->height;
vae_scale_factor = sd->get_vae_scale_factor();
diffusion_model_down_factor = sd->get_diffusion_model_down_factor();
seed = sd_img_gen_params->seed;
batch_count = sd_img_gen_params->batch_count;
qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers);
clip_skip = sd_img_gen_params->clip_skip;
shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep;
strength = sd_img_gen_params->strength;
control_strength = sd_img_gen_params->control_strength;
eta = sd_img_gen_params->sample_params.eta;
has_ref_images = sd_img_gen_params->ref_images_count > 0;
guidance = sd_img_gen_params->sample_params.guidance;
pm_params = sd_img_gen_params->pm_params;
pulid_params = sd_img_gen_params->pulid_params;
hires = sd_img_gen_params->hires;
cache_params = &sd_img_gen_params->cache;
resolve(sd);
}
GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params) {
prompt = SAFE_STR(sd_vid_gen_params->prompt);
negative_prompt = SAFE_STR(sd_vid_gen_params->negative_prompt);
width = sd_vid_gen_params->width;
height = sd_vid_gen_params->height;
requested_frames = std::max(1, sd_vid_gen_params->video_frames);
frames = sd->align_video_frames(requested_frames);
clip_skip = sd_vid_gen_params->clip_skip;
fps = std::max(1, sd_vid_gen_params->fps);
if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
fps = 24;
}
vae_scale_factor = sd->get_vae_scale_factor();
diffusion_model_down_factor = sd->get_diffusion_model_down_factor();
seed = sd_vid_gen_params->seed;
strength = sd_vid_gen_params->strength;
cache_params = &sd_vid_gen_params->cache;
vace_strength = sd_vid_gen_params->vace_strength;
guidance = sd_vid_gen_params->sample_params.guidance;
high_noise_guidance = sd_vid_gen_params->high_noise_sample_params.guidance;
hires = sd_vid_gen_params->hires;
resolve(sd);
if (frames != requested_frames) {
LOG_WARN("align video frames from %d to %d for %s",
requested_frames,
frames,
model_version_to_str[sd->version]);
}
}
void GenerationRequest::align_generation_request_size() {
align_image_size(&width, &height, "generation request");
}
void GenerationRequest::align_image_size(int* target_width, int* target_height, const char* label) {
int spatial_multiple = vae_scale_factor * diffusion_model_down_factor;
int width_offset = align_up_offset(*target_width, spatial_multiple);
int height_offset = align_up_offset(*target_height, spatial_multiple);
if (width_offset <= 0 && height_offset <= 0) {
return;
}
int original_width = *target_width;
int original_height = *target_height;
*target_width += width_offset;
*target_height += height_offset;
LOG_WARN("align %s up %dx%d to %dx%d (multiple=%d)",
label,
original_width,
original_height,
*target_width,
*target_height,
spatial_multiple);
}
void GenerationRequest::resolve_hires() {
if (!hires.enabled) {
return;
}
if (hires.upscaler == SD_HIRES_UPSCALER_NONE) {
hires.enabled = false;
return;
}
if (hires.upscaler < SD_HIRES_UPSCALER_NONE || hires.upscaler >= SD_HIRES_UPSCALER_COUNT) {
LOG_WARN("hires upscaler '%d' is invalid, disabling hires", hires.upscaler);
hires.enabled = false;
return;
}
if (hires.upscaler == SD_HIRES_UPSCALER_MODEL && strlen(SAFE_STR(hires.model_path)) == 0) {
LOG_WARN("hires model upscaler requires a model path, disabling hires");
hires.enabled = false;
return;
}
if (hires.scale <= 0.f && hires.target_width <= 0 && hires.target_height <= 0) {
LOG_WARN("hires scale must be positive when no target size is set, disabling hires");
hires.enabled = false;
return;
}
if (hires.custom_sigmas_count < 0) {
LOG_WARN("hires custom sigmas count is negative, ignoring custom sigmas");
hires.custom_sigmas = nullptr;
hires.custom_sigmas_count = 0;
}
if (hires.custom_sigmas_count > 0 && hires.custom_sigmas == nullptr) {
LOG_WARN("hires custom sigmas count is positive but custom sigmas are null, ignoring custom sigmas");
hires.custom_sigmas_count = 0;
}
if (hires.custom_sigmas_count == 1) {
LOG_WARN("hires custom sigmas requires at least two values, ignoring custom sigmas");
hires.custom_sigmas = nullptr;
hires.custom_sigmas_count = 0;
}
hires.denoising_strength = std::clamp(hires.denoising_strength, 0.0001f, 1.f);
hires.steps = std::max(0, hires.steps);
if (hires.target_width > 0 && hires.target_height > 0) {
// pass
} else if (hires.target_width > 0) {
hires.target_height = hires.target_width;
} else if (hires.target_height > 0) {
hires.target_width = hires.target_height;
} else {
hires.target_width = static_cast<int>(std::round(width * hires.scale));
hires.target_height = static_cast<int>(std::round(height * hires.scale));
}
if (hires.target_width <= 0 || hires.target_height <= 0) {
LOG_WARN("hires target size is not positive, disabling hires");
hires.enabled = false;
return;
}
align_image_size(&hires.target_width, &hires.target_height, "hires target");
}
void GenerationRequest::resolve_guidance(StableDiffusionGGML* sd,
sd_guidance_params_t* guidance,
bool* use_uncond,
bool* use_img_uncond,
bool has_ref_images,
const char* stage_name) {
GGML_ASSERT(guidance != nullptr);
GGML_ASSERT(use_uncond != nullptr);
GGML_ASSERT(use_img_uncond != nullptr);
// out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond)
// -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond
// out_cond : prompt, image latent
// out_uncond : negative prompt, image latent
// out_img_uncond : negative prompt, zero image latent
// image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG.
bool img_cfg_was_set = std::isfinite(guidance->img_cfg);
if (!img_cfg_was_set) {
guidance->img_cfg = 1.f;
}
if (!sd_version_supports_img_cfg(sd->version, has_ref_images)) {
if (img_cfg_was_set && guidance->img_cfg != 1.f) {
LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance");
}
guidance->img_cfg = 1.f;
}
if (guidance->img_cfg != guidance->txt_cfg) {
*use_uncond = true;
}
if (guidance->img_cfg != 1.f) {
*use_img_uncond = true;
}
if (guidance->txt_cfg < 1.f) {
const char* prefix = stage_name == nullptr ? "" : stage_name;
if (guidance->txt_cfg == 0.f) {
LOG_WARN("%sunconditioned mode, images won't follow the prompt (use cfg-scale=1 for distilled models)",
prefix);
} else {
LOG_WARN("%scfg value out of expected range may produce unexpected results", prefix);
}
}
}
void GenerationRequest::resolve(StableDiffusionGGML* sd) {
align_generation_request_size();
resolve_hires();
seed = resolve_seed(seed);
resolve_guidance(sd, &guidance, &use_uncond, &use_img_uncond, has_ref_images);
if (sd->high_noise_diffusion_model) {
resolve_guidance(sd,
&high_noise_guidance,
&use_high_noise_uncond,
&use_high_noise_img_uncond,
has_ref_images,
"high noise: ");
}
if (shifted_timestep > 0 && !sd_version_is_sdxl(sd->version)) {
LOG_WARN("timestep shifting is only supported for SDXL models!");
shifted_timestep = 0;
}
}
SamplePlan::SamplePlan(StableDiffusionGGML* sd,
const sd_img_gen_params_t* sd_img_gen_params,
const GenerationRequest& request) {
sample_method = sd_img_gen_params->sample_params.sample_method;
extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args;
eta = sd_img_gen_params->sample_params.eta;
sample_steps = sd_img_gen_params->sample_params.sample_steps;
resolve(sd, &request, &sd_img_gen_params->sample_params);
}
SamplePlan::SamplePlan(StableDiffusionGGML* sd,
const sd_vid_gen_params_t* sd_vid_gen_params,
const GenerationRequest& request) {
sample_method = sd_vid_gen_params->sample_params.sample_method;
extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args;
eta = sd_vid_gen_params->sample_params.eta;
sample_steps = sd_vid_gen_params->sample_params.sample_steps;
if (sd->high_noise_diffusion_model) {
high_noise_sample_steps = sd_vid_gen_params->high_noise_sample_params.sample_steps;
high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method;
high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args;
high_noise_eta = sd_vid_gen_params->high_noise_sample_params.eta;
}
moe_boundary = sd_vid_gen_params->moe_boundary;
resolve(sd, &request, &sd_vid_gen_params->sample_params);
}
void SamplePlan::resolve(StableDiffusionGGML* sd,
const GenerationRequest* request,
const sd_sample_params_t* sample_params) {
sample_method = resolve_sample_method(sd, sample_method);
total_steps = sample_steps + std::max(0, high_noise_sample_steps);
if (sample_params->custom_sigmas_count > 0) {
sigmas = std::vector<float>(sample_params->custom_sigmas,
sample_params->custom_sigmas + sample_params->custom_sigmas_count);
total_steps = static_cast<int>(sigmas.size()) - 1;
LOG_WARN("total_steps != custom_sigmas_count - 1, set total_steps to %d", total_steps);
if (sample_steps >= total_steps) {
sample_steps = total_steps;
LOG_WARN("total_steps != custom_sigmas_count - 1, set sample_steps to %d", sample_steps);
}
if (high_noise_sample_steps > 0) {
high_noise_sample_steps = total_steps - sample_steps;
LOG_WARN("total_steps != custom_sigmas_count - 1, set high_noise_sample_steps to %d", high_noise_sample_steps);
}
} else {
scheduler_t scheduler = resolve_scheduler(sd,
sample_params->scheduler,
sample_method);
int sample_seq_len = sd->get_image_seq_len(request->height, request->width);
if (sd_version_is_ltxav(sd->version) && request->frames > 0) {
int latent_frames = ((request->frames - 1) / 8) + 1;
sample_seq_len *= latent_frames;
} else if (sd_version_is_minimax_h3(sd->version) && request->frames > 0) {
sample_seq_len *= sd->video_frames_to_latent_frames(request->frames);
}
sigmas = sd->denoiser->get_sigmas(total_steps,
sample_seq_len,
scheduler,
sd->version,
sample_params->extra_sample_args);
}
eta = resolve_eta(sd, eta, sample_method);
if (high_noise_sample_steps < 0) {
for (size_t i = 0; i < sigmas.size(); ++i) {
if (sigmas[i] < moe_boundary) {
high_noise_sample_steps = static_cast<int>(i);
break;
}
}
LOG_VERBOSE("switching from high noise model at step %d", high_noise_sample_steps);
}
LOG_INFO("sampling using %s method", sampling_methods_str[sample_method]);
if (high_noise_sample_steps > 0) {
high_noise_sample_method = resolve_sample_method(sd,
high_noise_sample_method);
high_noise_eta = resolve_eta(sd, high_noise_eta, high_noise_sample_method);
LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]);
}
}
std::vector<float> make_hires_sigma_schedule(StableDiffusionGGML* sd,
const sd_hires_params_t& hires,
const sd_sample_params_t& sample_params,
sample_method_t sample_method,
int default_steps,
int sample_seq_len,
int* scheduler_steps_out) {
if (scheduler_steps_out != nullptr) {
*scheduler_steps_out = 0;
}
if (hires.custom_sigmas_count > 0 && hires.custom_sigmas != nullptr) {
std::vector<float> custom_sigmas(hires.custom_sigmas,
hires.custom_sigmas + hires.custom_sigmas_count);
if (scheduler_steps_out != nullptr) {
*scheduler_steps_out = static_cast<int>(custom_sigmas.size()) - 1;
}
return custom_sigmas;
}
int effective_steps = hires.steps > 0 ? hires.steps : default_steps;
effective_steps = std::max(1, effective_steps);
// sd-webui behavior: scale up total steps so trimming by denoising_strength yields exactly hires_steps effective steps,
// unlike img2img which trims from a fixed step count.
int scheduler_steps = static_cast<int>(effective_steps / hires.denoising_strength);
scheduler_steps = std::max(1, scheduler_steps);
scheduler_t scheduler = resolve_scheduler(sd,
sample_params.scheduler,
sample_method);
std::vector<float> sigmas = sd->denoiser->get_sigmas(scheduler_steps,
sample_seq_len,
scheduler,
sd->version,
sample_params.extra_sample_args);
size_t t_enc = static_cast<size_t>(scheduler_steps * hires.denoising_strength);
if (t_enc >= static_cast<size_t>(scheduler_steps)) {
t_enc = static_cast<size_t>(scheduler_steps) - 1;
}
if (scheduler_steps_out != nullptr) {
*scheduler_steps_out = scheduler_steps;
}
return std::vector<float>(sigmas.begin() + scheduler_steps - static_cast<int>(t_enc) - 1,
sigmas.end());
}
} // namespace sd::pipeline

110
src/pipeline/request.h Normal file
View File

@ -0,0 +1,110 @@
#ifndef __SD_PIPELINE_REQUEST_H__
#define __SD_PIPELINE_REQUEST_H__
#include <string>
#include <vector>
#include "stable-diffusion.h"
class StableDiffusionGGML;
namespace sd::pipeline {
extern const char* sampling_methods_str[];
enum sample_method_t default_sample_method(const StableDiffusionGGML* sd);
enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method);
float resolve_eta(StableDiffusionGGML* sd,
float eta,
enum sample_method_t sample_method);
struct GenerationRequest {
std::string prompt;
std::string negative_prompt;
int width = -1;
int height = -1;
int clip_skip = -1;
int vae_scale_factor = -1;
int diffusion_model_down_factor = -1;
int64_t seed = -1;
bool use_uncond = false;
bool use_img_uncond = false;
bool use_high_noise_uncond = false;
bool use_high_noise_img_uncond = false;
bool has_ref_images = false;
const sd_cache_params_t* cache_params = nullptr;
int batch_count = 1;
int qwen_image_layers = 3;
int shifted_timestep = 0;
float strength = 1.f;
float control_strength = 0.f;
float eta = 0.f;
sd_guidance_params_t guidance = {};
sd_guidance_params_t high_noise_guidance = {};
sd_pm_params_t pm_params = {};
sd_pulid_params_t pulid_params = {};
sd_hires_params_t hires = {};
int frames = -1;
int requested_frames = -1;
int fps = 16;
float vace_strength = 1.f;
GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params);
GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params);
void align_generation_request_size();
void align_image_size(int* target_width, int* target_height, const char* label);
void resolve_hires();
static void resolve_guidance(StableDiffusionGGML* sd,
sd_guidance_params_t* guidance,
bool* use_uncond,
bool* use_img_uncond,
bool has_ref_images,
const char* stage_name = nullptr);
void resolve(StableDiffusionGGML* sd);
};
struct SamplePlan {
enum sample_method_t sample_method = SAMPLE_METHOD_COUNT;
enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT;
const char* extra_sample_args = nullptr;
const char* high_noise_extra_sample_args = nullptr;
float eta = 0.f;
float high_noise_eta = 0.f;
int sample_steps = 0;
int high_noise_sample_steps = 0;
int total_steps = 0;
float moe_boundary = 0.f;
std::vector<float> sigmas;
SamplePlan(StableDiffusionGGML* sd,
const sd_img_gen_params_t* sd_img_gen_params,
const GenerationRequest& request);
SamplePlan(StableDiffusionGGML* sd,
const sd_vid_gen_params_t* sd_vid_gen_params,
const GenerationRequest& request);
void resolve(StableDiffusionGGML* sd,
const GenerationRequest* request,
const sd_sample_params_t* sample_params);
};
std::vector<float> make_hires_sigma_schedule(StableDiffusionGGML* sd,
const sd_hires_params_t& hires,
const sd_sample_params_t& sample_params,
sample_method_t sample_method,
int default_steps,
int sample_seq_len,
int* scheduler_steps_out);
} // namespace sd::pipeline
#endif // __SD_PIPELINE_REQUEST_H__

1797
src/pipeline/video.cpp Normal file

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

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