From bbc29e632cd458f40e603ebbbddc19e5603ccc9b Mon Sep 17 00:00:00 2001 From: fszontagh Date: Fri, 10 Jul 2026 11:22:14 +0200 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20hot-reload=20ControlNet=20=E2=80=94?= =?UTF-8?q?=20swap=20without=20rebuilding=20the=20context?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three new C API entry points on stable-diffusion.h so a long-lived process (server / GUI / integration) can swap the ControlNet weights without paying the cost of tearing down the ctx and reloading the base diffusion / VAE / text encoders (multi-GB, tens of seconds): SD_API bool sd_ctx_load_control_net(sd_ctx_t*, const char* path); SD_API bool sd_ctx_unload_control_net(sd_ctx_t*); SD_API bool sd_ctx_has_control_net(const sd_ctx_t*); Implementation: - StableDiffusionGGML gains load_control_net_from_file(path) and unload_control_net() helpers. They reuse the existing ControlNet construction path (backend pair, conv2d-direct setting, register_runner_params), so hot- loaded ControlNets are indistinguishable at runtime from init-time ones. - ModelManager gains unregister_param_tensors(desc) which releases both the compute and params backend buffers for a runner's tensors and removes them from tensor_states_by_name_ / tensor_states_, so a subsequent register with the same tensor names does not collide. - diffusion_conv_direct + control_net_params_mem_size are promoted from local-scope in the initial-load path to StableDiffusionGGML members so that the reload path can honor the ctx's original conv-direct setting and keep the memory bookkeeping consistent. - The three C API functions are thin wrappers over those helpers, matching the existing sd_ctx_* pattern in stable-diffusion.h. Doc caveats in the header note that this API is intended for use when no generation is in flight — thread safety is the caller's responsibility, same as the rest of the sd_ctx API surface. Smoke test (create ctx w/o ControlNet, load canny, swap to depth, unload, reload canny, free) exits cleanly with correct has_control_net transitions. Existing --control-net init-time path is unchanged (regression-tested with sd-cli). Motivation is to unblock two follow-ups: 1. Servers can offer a ControlNet dropdown without a full model reload. 2. The same pattern generalizes to other module families (LoRA, VAE, upscaler); the private helpers were factored with that extension in mind. --- include/stable-diffusion.h | 14 ++++++ src/model_manager.cpp | 49 ++++++++++++++++++++ src/model_manager.h | 3 ++ src/stable-diffusion.cpp | 95 +++++++++++++++++++++++++++++++++++++- 4 files changed, 160 insertions(+), 1 deletion(-) diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 93084fe29..5cde63681 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -428,6 +428,20 @@ SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx); +// Hot-swap the ControlNet weights on an already-initialized context. +// Loads the file at `path`, replaces any currently loaded ControlNet, and re- +// registers it with the model manager. The file must be compatible with the +// diffusion model the context was created with. Safe to call only when no +// generation is in flight. Returns true on success. +SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path); + +// Release the currently loaded ControlNet (if any). Subsequent generations +// will run without ControlNet guidance. Always returns true. +SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx); + +// Returns true if a ControlNet is currently loaded on the context. +SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx); + SD_API const char* sd_type_name(enum sd_type_t type); SD_API enum sd_type_t str_to_sd_type(const char* str); SD_API const char* sd_rng_type_name(enum rng_type_t rng_type); diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 3a98bd545..013bb3a4b 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -179,6 +179,55 @@ bool ModelManager::register_param_tensors(const std::string& desc, return true; } +void ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) { + if (desc.empty()) { + return; + } + + std::unordered_set target_states; + std::vector backend_tensors; + size_t released_size = 0; + for (auto& state : tensor_states_) { + if (state == nullptr || state->desc != desc) { + continue; + } + target_states.insert(state.get()); + if (state->tensor != nullptr) { + backend_tensors.push_back(state->tensor); + released_size += ggml_nbytes(state->tensor); + } + } + + if (target_states.empty()) { + return; + } + + release_compute_backend_params(backend_tensors); + release_params_backend_params(backend_tensors); + + for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) { + if (target_states.count(it->second) > 0) { + it = tensor_states_by_name_.erase(it); + } else { + ++it; + } + } + tensor_states_.erase(std::remove_if(tensor_states_.begin(), + tensor_states_.end(), + [&](const std::unique_ptr& s) { + return s == nullptr || target_states.count(s.get()) > 0; + }), + tensor_states_.end()); + + if (registered_tensor_size != nullptr) { + if (released_size > *registered_tensor_size) { + *registered_tensor_size = 0; + } else { + *registered_tensor_size -= released_size; + } + } +} + bool ModelManager::load_all_params_eagerly() { std::vector all_states; all_states.reserve(tensor_states_.size()); diff --git a/src/model_manager.h b/src/model_manager.h index d80032614..baa0631b4 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -134,6 +134,9 @@ class ModelManager : public RunnerWeightManager { bool allow_split_buffer = false, bool params_follow_compute_backend = false); + void unregister_param_tensors(const std::string& desc, + size_t* registered_tensor_size = nullptr); + template bool register_runner_params(const std::string& desc, Runner& runner, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 28628be7c..7b8bdd937 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -222,9 +222,13 @@ class StableDiffusionGGML { 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; + size_t control_net_params_mem_size = 0; + std::shared_ptr model_manager; std::shared_ptr denoiser = std::make_shared(); @@ -494,6 +498,72 @@ class StableDiffusionGGML { params_follow_runtime); } + bool unload_control_net() { + if (control_net == nullptr) { + return true; + } + if (model_manager != nullptr) { + model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size); + } + control_net.reset(); + control_net_params_mem_size = 0; + return true; + } + + bool load_control_net_from_file(const std::string& path) { + if (path.empty()) { + LOG_ERROR("sd_ctx_load_control_net: empty path"); + return false; + } + if (model_manager == nullptr) { + LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized"); + return false; + } + + unload_control_net(); + + ModelLoader& shared_loader = model_manager->loader(); + if (!shared_loader.init_from_file(path)) { + LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str()); + return false; + } + shared_loader.convert_tensors_name(); + + if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) { + LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable"); + return false; + } + + control_net = std::make_shared(backend_for(SDBackendModule::CONTROL_NET), + params_backend_for(SDBackendModule::CONTROL_NET), + shared_loader.get_tensor_storage_map(), + version, + "", + model_manager); + if (diffusion_conv_direct) { + LOG_INFO("Using Conv2d direct in the control net"); + control_net->set_conv2d_direct_enabled(true); + } + if (!register_runner_params("ControlNet", + control_net, + SDBackendModule::CONTROL_NET, + &control_net_params_mem_size)) { + LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed"); + control_net.reset(); + control_net_params_mem_size = 0; + return false; + } + if (!model_manager->validate_registered_tensors()) { + LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed"); + unload_control_net(); + return false; + } + LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)", + path.c_str(), + control_net_params_mem_size / 1024.0 / 1024.0); + return true; + } + bool init_backend() { std::string error; if (!backend_manager.init(backend_spec.c_str(), @@ -852,10 +922,12 @@ class StableDiffusionGGML { model_loader.process_model_files(enable_mmap, needs_writable_mmap); load_alphas_cumprod(model_loader); + diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; + size_t text_encoder_params_mem_size = 0; size_t unet_params_mem_size = 0; size_t vae_params_mem_size = 0; - size_t control_net_params_mem_size = 0; + control_net_params_mem_size = 0; size_t extension_params_mem_size = 0; bool tae_preview_only = sd_ctx_params->tae_preview_only; @@ -3429,6 +3501,27 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { return sd_version_supports_video_generation(sd_ctx->sd->version); } +SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { + return false; + } + return sd_ctx->sd->load_control_net_from_file(path); +} + +SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; + } + return sd_ctx->sd->unload_control_net(); +} + +SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; + } + return sd_ctx->sd->control_net != nullptr; +} + enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { if (sd_ctx != nullptr && sd_ctx->sd != nullptr) { if (sd_version_is_pid(sd_ctx->sd->version)) { From 3d8b90aa241845252c124a763146402703faffa0 Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 10 Jul 2026 23:04:31 +0800 Subject: [PATCH 2/3] fix: safely unregister hot-swapped ControlNet tensors --- src/model_manager.cpp | 61 +++++++++++++++++++++++++++++++++++----- src/model_manager.h | 2 +- src/stable-diffusion.cpp | 8 ++++-- 3 files changed, 61 insertions(+), 10 deletions(-) diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 013bb3a4b..33bc18373 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -179,31 +179,77 @@ bool ModelManager::register_param_tensors(const std::string& desc, return true; } -void ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) { +bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) { if (desc.empty()) { - return; + return true; } std::unordered_set target_states; - std::vector backend_tensors; size_t released_size = 0; for (auto& state : tensor_states_) { if (state == nullptr || state->desc != desc) { continue; } + if (state->active_prepare_count > 0) { + LOG_ERROR("model manager cannot unregister active %s tensor '%s'", + desc.c_str(), + state->name.c_str()); + return false; + } target_states.insert(state.get()); if (state->tensor != nullptr) { - backend_tensors.push_back(state->tensor); released_size += ggml_nbytes(state->tensor); } } if (target_states.empty()) { - return; + return true; } - release_compute_backend_params(backend_tensors); - release_params_backend_params(backend_tensors); + release_compute_staging_blocks(false); + + std::vector storage_blocks_to_release; + std::unordered_set affected_storage_states; + for (const auto& block : params_storage_blocks_) { + if (block == nullptr) { + continue; + } + bool has_target_state = false; + for (TensorState* state : block->states) { + if (state != nullptr && target_states.count(state) > 0) { + has_target_state = true; + break; + } + } + if (!has_target_state) { + continue; + } + storage_blocks_to_release.push_back(block.get()); + for (TensorState* state : block->states) { + if (state != nullptr) { + affected_storage_states.insert(state); + } + } + } + + for (TensorState* state : affected_storage_states) { + if (state == nullptr) { + continue; + } + if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active", + desc.c_str(), + state->name.c_str()); + return false; + } + } + + for (ParamsStorageBlock* block : storage_blocks_to_release) { + if (block != nullptr) { + free_params_storage_block(*block); + erase_params_storage_block(block); + } + } for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) { if (target_states.count(it->second) > 0) { @@ -226,6 +272,7 @@ void ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg *registered_tensor_size -= released_size; } } + return true; } bool ModelManager::load_all_params_eagerly() { diff --git a/src/model_manager.h b/src/model_manager.h index baa0631b4..85e982de2 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -134,7 +134,7 @@ class ModelManager : public RunnerWeightManager { bool allow_split_buffer = false, bool params_follow_compute_backend = false); - void unregister_param_tensors(const std::string& desc, + bool unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size = nullptr); template diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 7b8bdd937..91edaced8 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -503,7 +503,9 @@ class StableDiffusionGGML { return true; } if (model_manager != nullptr) { - model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size); + if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) { + return false; + } } control_net.reset(); control_net_params_mem_size = 0; @@ -520,7 +522,9 @@ class StableDiffusionGGML { return false; } - unload_control_net(); + if (!unload_control_net()) { + return false; + } ModelLoader& shared_loader = model_manager->loader(); if (!shared_loader.init_from_file(path)) { From 020858bb049c15c293d4b793813dae6c35e758e1 Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 10 Jul 2026 23:06:30 +0800 Subject: [PATCH 3/3] clean code --- include/stable-diffusion.h | 11 +---------- src/stable-diffusion.cpp | 2 +- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 5cde63681..eeb59f873 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -428,18 +428,9 @@ SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx); -// Hot-swap the ControlNet weights on an already-initialized context. -// Loads the file at `path`, replaces any currently loaded ControlNet, and re- -// registers it with the model manager. The file must be compatible with the -// diffusion model the context was created with. Safe to call only when no -// generation is in flight. Returns true on success. +// ControlNet hot-swap APIs are not safe to call while generation is in flight. SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path); - -// Release the currently loaded ControlNet (if any). Subsequent generations -// will run without ControlNet guidance. Always returns true. SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx); - -// Returns true if a ControlNet is currently loaded on the context. SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx); SD_API const char* sd_type_name(enum sd_type_t type); diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 91edaced8..1c1563b52 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -926,7 +926,7 @@ class StableDiffusionGGML { model_loader.process_model_files(enable_mmap, needs_writable_mmap); load_alphas_cumprod(model_loader); - diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; + diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; size_t text_encoder_params_mem_size = 0; size_t unet_params_mem_size = 0;