feat: support disk params backend (#1651)

This commit is contained in:
leejet authored and GitHub committed 2026-06-14 14:48:50 +08:00
1 parent 276025e054
commit bdb431ad95
27 files changed
+134 -65

No files matched your search

+28 -10
View File
@@ -45,6 +45,10 @@ static bool is_default_backend_token(const std::string& name) {
return lower.empty() || lower == "default" || lower == "auto";
}
static bool is_disk_backend_token(const std::string& name) {
return lower_copy(trim_copy(name)) == "disk";
}
static bool parse_backend_module(const std::string& raw_name, SDBackendModule* module) {
std::string name = lower_copy(trim_copy(raw_name));
name.erase(std::remove(name.begin(), name.end(), '-'), name.end());
@@ -504,6 +508,9 @@ ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
if (name.empty()) {
return runtime_backend(module);
}
if (is_disk_backend_token(name)) {
return runtime_backend(module);
}
return init_cached_backend(name);
}
@@ -515,6 +522,10 @@ bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
return sd_backend_is_cpu(params_backend(module));
}
bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const {
return is_disk_backend_token(params_assignment_.get(module));
}
bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
ggml_backend_t backend = runtime_backend(module);
if (backend == nullptr) {
@@ -534,7 +545,6 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
bool SDBackendManager::init(const char* backend_spec,
const char* params_backend_spec,
bool offload_params_to_cpu,
bool keep_clip_on_cpu,
bool keep_vae_on_cpu,
bool keep_control_net_on_cpu,
@@ -560,18 +570,20 @@ bool SDBackendManager::init(const char* backend_spec,
}
}
if (params_assignment_.empty() && offload_params_to_cpu) {
params_assignment_.set_default("cpu");
}
return validate(error);
}
bool SDBackendManager::validate(std::string* error) const {
auto validate_name = [&](const std::string& name) -> bool {
auto validate_runtime_name = [&](const std::string& name) -> bool {
if (is_default_backend_token(name)) {
return true;
}
if (is_disk_backend_token(name)) {
if (error != nullptr) {
*error = "backend 'disk' is only supported by params_backend";
}
return false;
}
if (!sd_resolve_backend_name(name).empty()) {
return true;
}
@@ -580,18 +592,24 @@ bool SDBackendManager::validate(std::string* error) const {
}
return false;
};
auto validate_params_name = [&](const std::string& name) -> bool {
if (is_disk_backend_token(name)) {
return true;
}
return validate_runtime_name(name);
};
if (!validate_name(runtime_assignment_.default_name) ||
!validate_name(params_assignment_.default_name)) {
if (!validate_runtime_name(runtime_assignment_.default_name) ||
!validate_params_name(params_assignment_.default_name)) {
return false;
}
for (const auto& kv : runtime_assignment_.module_names) {
if (!validate_name(kv.second)) {
if (!validate_runtime_name(kv.second)) {
return false;
}
}
for (const auto& kv : params_assignment_.module_names) {
if (!validate_name(kv.second)) {
if (!validate_params_name(kv.second)) {
return false;
}
}
+1 -1
View File
@@ -51,7 +51,6 @@ public:
bool init(const char* backend_spec,
const char* params_backend_spec,
bool offload_params_to_cpu,
bool keep_clip_on_cpu,
bool keep_vae_on_cpu,
bool keep_control_net_on_cpu,
@@ -63,6 +62,7 @@ public:
bool runtime_backend_is_cpu(SDBackendModule module);
bool params_backend_is_cpu(SDBackendModule module);
bool params_backend_is_disk(SDBackendModule module) const;
bool runtime_backend_supports_host_buffer(SDBackendModule module);
private:
+1 -1
View File
@@ -101,7 +101,7 @@ struct LoraModel : public GGMLRunner {
if (model_manager == nullptr ||
!model_manager->register_param_tensors("LoRA",
std::move(tensors),
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -622,7 +622,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
model_loader.load_tensors(on_new_tensor_cb);
if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
tensors,
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -482,7 +482,7 @@ struct ControlNet : public GGMLRunner {
manager->set_n_threads(n_threads);
if (!manager->register_param_tensors("ControlNet",
std::move(tensors),
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
runtime_backend,
params_backend) ||
!manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1609,7 +1609,7 @@ namespace Flux {
if (!model_manager->register_runner_params("Flux test",
*flux,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -2048,7 +2048,7 @@ namespace LTXV {
if (!model_manager->register_runner_params("LTXAV test",
*ltxav,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1015,7 +1015,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
if (!model_manager->register_runner_params("MMDiT test",
*mmdit,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -715,7 +715,7 @@ namespace Qwen {
if (!model_manager->register_runner_params("Qwen image test",
*qwen_image,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1040,7 +1040,7 @@ namespace WAN {
if (!model_manager->register_runner_params("Wan test",
*wan,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -723,7 +723,7 @@ namespace ZImage {
if (!model_manager->register_runner_params("ZImage test",
*z_image,
"model.diffusion_model",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -2084,7 +2084,7 @@ namespace LLM {
if (!model_manager->register_runner_params("LLM test",
*llm,
"text_encoders.llm",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -592,7 +592,7 @@ struct T5Embedder {
if (!model_manager->register_runner_params("T5 test",
*t5,
"",
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1082,7 +1082,7 @@ namespace LTXV {
if (!model_manager->register_runner_params("LTX audio VAE test",
*ltx_audio_vae,
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1538,7 +1538,7 @@ struct LTXVideoVAE : public VAE {
if (!model_manager->register_runner_params("LTX VAE test",
*vae,
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -1340,7 +1340,7 @@ namespace WAN {
if (!model_manager->register_runner_params("Wan VAE test",
*vae,
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
backend,
backend) ||
!model_manager->validate_registered_tensors()) {
+1 -1
View File
@@ -492,7 +492,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
}
bool ModelManager::can_mmap_storage(const TensorState& state) const {
if (!enable_mmap_ || state.residency_mode != ResidencyMode::Resident) {
if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
return false;
}
if (state.compute_backend == nullptr || state.params_backend == nullptr) {
+2 -2
View File
@@ -16,7 +16,7 @@ class ModelManager : public RunnerWeightManager {
public:
enum class ResidencyMode {
Disk,
Resident,
ParamBackend,
};
struct LoraSpec {
@@ -33,7 +33,7 @@ private:
ggml_tensor* tensor = nullptr;
std::string desc;
ResidencyMode residency_mode = ResidencyMode::Resident;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
bool metadata_validated = false;
+18 -22
View File
@@ -165,7 +165,6 @@ public:
SDVersion version;
bool vae_decode_only = false;
bool external_vae_is_invalid = false;
bool free_params_immediately = false;
bool circular_x = false;
bool circular_y = false;
@@ -246,7 +245,7 @@ public:
}
return model_manager->register_param_tensors(desc,
std::move(group_tensors),
free_params_immediately ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::Resident,
backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
backend_for(module),
params_backend_for(module),
params_mem_size);
@@ -255,8 +254,7 @@ public:
bool init_backend(const sd_ctx_params_t* sd_ctx_params) {
std::string error;
if (!backend_manager.init(sd_ctx_params->backend,
sd_ctx_params->params_backend,
offload_params_to_cpu,
params_backend_spec.c_str(),
sd_ctx_params->keep_clip_on_cpu,
sd_ctx_params->keep_vae_on_cpu,
sd_ctx_params->keep_control_net_on_cpu,
@@ -319,24 +317,21 @@ public:
}
bool init(const sd_ctx_params_t* sd_ctx_params) {
n_threads = sd_ctx_params->n_threads;
vae_decode_only = sd_ctx_params->vae_decode_only;
free_params_immediately = sd_ctx_params->free_params_immediately;
offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;
enable_mmap = sd_ctx_params->enable_mmap;
max_vram = sd_ctx_params->max_vram;
stream_layers = sd_ctx_params->stream_layers;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
n_threads = sd_ctx_params->n_threads;
vae_decode_only = sd_ctx_params->vae_decode_only;
offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;
enable_mmap = sd_ctx_params->enable_mmap;
max_vram = sd_ctx_params->max_vram;
stream_layers = sd_ctx_params->stream_layers;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
if (offload_params_to_cpu) {
params_backend_spec = params_backend_spec.empty() ? "*=cpu" : "*=cpu," + params_backend_spec;
}
if (stream_layers && max_vram == 0.f) {
LOG_WARN("--stream-layers has no effect without --max-vram set; ignoring");
stream_layers = false;
}
if (stream_layers && !offload_params_to_cpu && params_backend_spec.empty()) {
// Streaming needs CPU-resident params.
LOG_WARN("--stream-layers has no effect without --offload-to-cpu (or --params-backend); ignoring");
stream_layers = false;
}
bool use_tae = false;
bool use_audio_vae = false;
@@ -354,6 +349,10 @@ public:
if (!init_backend(sd_ctx_params)) {
return false;
}
if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false;
}
max_vram = sd::ggml_graph_cut::resolve_max_vram_gib(max_vram, backend_for(SDBackendModule::DIFFUSION));
model_manager = std::make_shared<ModelManager>();
@@ -2644,7 +2643,6 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) {
void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
*sd_ctx_params = {};
sd_ctx_params->vae_decode_only = true;
sd_ctx_params->free_params_immediately = true;
sd_ctx_params->n_threads = sd_get_num_physical_cores();
sd_ctx_params->wtype = SD_TYPE_COUNT;
sd_ctx_params->rng_type = CUDA_RNG;
@@ -2694,7 +2692,6 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"photo_maker_path: %s\n"
"tensor_type_rules: %s\n"
"vae_decode_only: %s\n"
"free_params_immediately: %s\n"
"n_threads: %d\n"
"wtype: %s\n"
"rng_type: %s\n"
@@ -2734,7 +2731,6 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
SAFE_STR(sd_ctx_params->photo_maker_path),
SAFE_STR(sd_ctx_params->tensor_type_rules),
BOOL_STR(sd_ctx_params->vae_decode_only),
BOOL_STR(sd_ctx_params->free_params_immediately),
sd_ctx_params->n_threads,
sd_type_name(sd_ctx_params->wtype),
sd_rng_type_name(sd_ctx_params->rng_type),
@@ -5037,7 +5033,7 @@ static sd::Tensor<float> upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx,
upsampler->get_param_tensors(tensors);
if (!upsampler_manager->register_param_tensors("LTX latent upsampler",
std::move(tensors),
ModelManager::ResidencyMode::Resident,
ModelManager::ResidencyMode::ParamBackend,
sd_ctx->sd->backend_for(SDBackendModule::UPSCALER),
sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)) ||
!upsampler_manager->validate_registered_tensors()) {
+6 -3
View File
@@ -43,10 +43,13 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
int n_threads) {
ggml_log_set(ggml_log_callback_default, nullptr);
std::string effective_params_backend_spec = params_backend_spec;
if (offload_params_to_cpu) {
effective_params_backend_spec = effective_params_backend_spec.empty() ? "*=cpu" : "*=cpu," + effective_params_backend_spec;
}
std::string error;
if (!backend_manager.init(backend_spec.c_str(),
params_backend_spec.c_str(),
offload_params_to_cpu,
effective_params_backend_spec.c_str(),
false,
false,
false,
@@ -106,7 +109,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
esrgan_upscaler->get_param_tensors(tensors);
if (!model_manager->register_param_tensors("ESRGAN",
std::move(tensors),
ModelManager::ResidencyMode::Resident,
backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
backend_for(SDBackendModule::UPSCALER),
params_backend_for(SDBackendModule::UPSCALER)) ||
!model_manager->validate_registered_tensors()) {