diff --git a/src/abot_world.hpp b/src/abot_world.hpp index 605cb3fef..b24152826 100644 --- a/src/abot_world.hpp +++ b/src/abot_world.hpp @@ -974,9 +974,16 @@ class AbotWalkSession { // return false to fall back to the RNG. std::function noise_override; - std::shared_ptr model_manager; std::unique_ptr runner; std::shared_ptr tae; + // Declared after the runners on purpose (mirrors StableDiffusionGGML): + // ~ModelManager force-frees the param storage blocks and writes through + // the registered ggml tensors (state->tensor->buffer = nullptr), which + // live in the runner/tae contexts above. Members destroy in reverse + // declaration order and the runners hold only weak_ptr refs to the + // manager, so this ordering runs ~ModelManager first, while every + // registered tensor is still alive. + std::shared_ptr model_manager; // finalized walk state (per-frame ggml {W,H,C} latents, torch [C,H,W] flat) std::vector> history; diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index 017d3d8f3..784683f29 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -2071,7 +2071,12 @@ struct GGMLRunner { if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) { cpu_fallback_backend = sd_backend_cpu_init(); } - if (cpu_fallback_backend != nullptr) { + // The internal CPU fallback backend is only needed as the + // scheduler's trailing fallback while the primary runtime is + // non-CPU. When a VAE graph is retried on its explicit CPU backend, + // adding a second CPU backend creates an ambiguous scheduler route. + if (cpu_fallback_backend != nullptr && + !sd_backend_is_cpu(runtime_backend)) { backends.push_back(cpu_fallback_backend); } @@ -2227,6 +2232,36 @@ struct GGMLRunner { return bytes; } + bool assign_graph_params_compute_backend(const std::vector& graph_params, + ggml_backend_t backend, + const char* action) { + GGML_ASSERT(backend != nullptr); + auto manager = weight_manager.lock(); + if (manager == nullptr) { + if (graph_params.empty()) { + return true; + } + LOG_ERROR("%s VAE CPU fallback cannot %s graph params without a weight manager", + get_desc().c_str(), + action); + return false; + } + if (!manager->assign_compute_backend(graph_params, backend)) { + LOG_ERROR("%s VAE CPU fallback failed to %s %zu graph params to %s", + get_desc().c_str(), + action, + graph_params.size(), + ggml_backend_name(backend)); + return false; + } + LOG_DEBUG("%s VAE CPU fallback %s %zu graph params to %s", + get_desc().c_str(), + action, + graph_params.size(), + ggml_backend_name(backend)); + return true; + } + void switch_runtime_backend(ggml_backend_t backend) { GGML_ASSERT(backend != nullptr); free_compute_buffer(); @@ -3233,37 +3268,59 @@ struct GGMLRunner { } const bool has_stateful_cache = !cache_tensor_map.empty() || cache_ctx != nullptr; - auto retry_stateless_on_cpu = - [&](const char* failure) -> std::optional> { - if (!vae_auto_cpu_fallback_enabled || - sd_backend_is_cpu(runtime_backend) || - vae_fallback_backend == nullptr || - !sd_backend_is_cpu(vae_fallback_backend) || - has_stateful_cache) { + auto run_stateless_on_cpu = + [&](const std::vector& graph_params) -> std::optional> { + ggml_backend_t previous_backend = runtime_backend; + const std::string previous_backend_name = ggml_backend_name(previous_backend); + if (!assign_graph_params_compute_backend(graph_params, + vae_fallback_backend, + "assign")) { return std::nullopt; } - ggml_backend_t previous_backend = runtime_backend; - const std::string previous_backend_name = - ggml_backend_name(previous_backend); - LOG_WARN("%s VAE %s on %s; retrying stateless graph on CPU", - get_desc().c_str(), - failure, - previous_backend_name.c_str()); switch_runtime_backend(vae_fallback_backend); + + auto restore_previous_backend = [&]() { + const bool restored = assign_graph_params_compute_backend(graph_params, + previous_backend, + "restore"); + switch_runtime_backend(previous_backend); + return restored; + }; + try { + // Fallback staging cannot remain active after runtime restoration. + // Always release CPU compute params before assigning the graph back. auto output = compute(get_graph, n_threads, false, free_compute_buffer, - free_compute_params, + true, no_return); - switch_runtime_backend(previous_backend); + if (!restore_previous_backend()) { + return std::nullopt; + } return output; } catch (...) { - switch_runtime_backend(previous_backend); + restore_previous_backend(); throw; } }; + auto retry_stateless_on_cpu = + [&](const char* failure) -> std::optional> { + if (!vae_auto_cpu_fallback_enabled || + sd_backend_is_cpu(runtime_backend) || + vae_fallback_backend == nullptr || + !sd_backend_is_cpu(vae_fallback_backend) || + has_stateful_cache) { + return std::nullopt; + } + const std::string previous_backend_name = ggml_backend_name(runtime_backend); + LOG_WARN("%s VAE %s on %s; retrying stateless graph on CPU", + get_desc().c_str(), + failure, + previous_backend_name.c_str()); + return run_stateless_on_cpu(collect_used_param_tensors(gf)); + }; if (vae_auto_cpu_fallback_enabled && !sd_backend_is_cpu(runtime_backend) && @@ -3315,23 +3372,13 @@ struct GGMLRunner { previous_backend_name.c_str(), cpu_backend_name.c_str()); - switch_runtime_backend(vae_fallback_backend); - try { - auto output = compute(get_graph, - n_threads, - false, - free_compute_buffer, - free_compute_params, - no_return); - switch_runtime_backend(previous_backend); + auto output = run_stateless_on_cpu(collect_used_param_tensors(gf)); + if (output.has_value()) { LOG_INFO("%s VAE CPU fallback complete; restored runtime backend %s", get_desc().c_str(), previous_backend_name.c_str()); - return output; - } catch (...) { - switch_runtime_backend(previous_backend); - throw; } + return output; } if (decision.reason == sd::VaeGraphRouteReason::STATEFUL_GRAPH) { diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index e89ecf492..5005af2e1 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -7517,16 +7517,23 @@ bool sd_abot_scene_create(const sd_abot_scene_params_t* p) { LOG_ERROR("sd_abot_scene_create: backend init failed: %s", error.c_str()); return false; } const int threads = p->n_threads > 0 ? p->n_threads : sd_get_num_physical_cores(); + // The runners below own the ggml tensors that ~ModelManager's + // storage-block teardown writes through (state->tensor->buffer = + // nullptr), and they hold only weak_ptr refs to the manager. Declare + // them BEFORE the manager so reverse local destruction runs + // ~ModelManager first, while every registered tensor is still alive + // (mirrors the member ordering in StableDiffusionGGML). + std::shared_ptr t5; + std::shared_ptr vae; auto model_manager = std::make_shared(); model_manager->set_n_threads(threads); ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file(p->t5_path, "text_encoders.t5xxl.transformer.")) return false; - auto t5 = std::make_shared(backends.runtime_backend(SDBackendModule::TE), - model_loader.get_tensor_storage_map(), - "text_encoders.t5xxl.transformer", - true, - model_manager); - std::shared_ptr vae; + t5 = std::make_shared(backends.runtime_backend(SDBackendModule::TE), + model_loader.get_tensor_storage_map(), + "text_encoders.t5xxl.transformer", + true, + model_manager); if (has_image) { if (!model_loader.init_from_file(p->vae_path, "first_stage_model.")) return false; vae = std::make_shared(backends.runtime_backend(SDBackendModule::VAE), diff --git a/tests/test-vae-routing.cpp b/tests/test-vae-routing.cpp index eb5d8b5b5..2e9314899 100644 --- a/tests/test-vae-routing.cpp +++ b/tests/test-vae-routing.cpp @@ -1,12 +1,53 @@ +#include "core/ggml_extend.hpp" #include "core/ggml_extend_backend.h" #include "vae_fallback.hpp" #include #include +#include #include +#include namespace { constexpr size_t MIB = 1024ull * 1024ull; + + class RecordingWeightManager : public RunnerWeightManager { + public: + std::vector assignments; + size_t fail_on_assignment = 0; + + bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) override { + GGML_ASSERT(!tensors.empty()); + assignments.push_back(compute_backend); + return fail_on_assignment == 0 || assignments.size() != fail_on_assignment; + } + + bool prepare_params(const std::vector&) override { return true; } + void release_compute_backend_params(const std::vector&) override {} + void release_params_backend_params(const std::vector&) override {} + }; + + class RoutingTestRunner : public GGMLRunner { + public: + RoutingTestRunner(ggml_backend_t backend, + const std::shared_ptr& manager) + : GGMLRunner(backend, manager) {} + + std::string get_desc() override { return "routing_test"; } + + ggml_tensor* make_param() { + ggml_tensor* tensor = ggml_new_tensor_1d(params_ctx, GGML_TYPE_F32, 1); + ggml_set_name(tensor, "routing_test.weight"); + return tensor; + } + + bool assign_graph_params(const std::vector& tensors, + ggml_backend_t backend, + const char* action) { + return assign_graph_params_compute_backend(tensors, backend, action); + } + }; } int main() { @@ -21,6 +62,30 @@ int main() { &compatibility_error)); compatibility_manager.reset(); + ggml_backend_t original_backend = sd_backend_cpu_init(); + ggml_backend_t fallback_backend = sd_backend_cpu_init(); + GGML_ASSERT(original_backend != nullptr); + GGML_ASSERT(fallback_backend != nullptr); + { + auto manager = std::make_shared(); + RoutingTestRunner runner(original_backend, manager); + std::vector graph_params = {runner.make_param()}; + + GGML_ASSERT(runner.assign_graph_params(graph_params, fallback_backend, "assign")); + GGML_ASSERT(runner.assign_graph_params(graph_params, original_backend, "restore")); + GGML_ASSERT(manager->assignments.size() == 2); + GGML_ASSERT(manager->assignments[0] == fallback_backend); + GGML_ASSERT(manager->assignments[1] == original_backend); + + manager->assignments.clear(); + manager->fail_on_assignment = 2; + GGML_ASSERT(runner.assign_graph_params(graph_params, fallback_backend, "assign")); + GGML_ASSERT(!runner.assign_graph_params(graph_params, original_backend, "restore")); + GGML_ASSERT(manager->assignments.size() == 2); + } + ggml_backend_free(fallback_backend); + ggml_backend_free(original_backend); + SDBackendAssignment runtime; runtime.set_default("vulkan0"); runtime.set_module(SDBackendModule::VAE, "metal");