refactor: centralize runner weight staging and cleanup (#1644)

src/model/diffusion/anima.hpp

@@ -697,7 +697,7 @@ namespace Anima {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/control.hpp

@@ -309,6 +309,7 @@ public:
 struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
     ControlNetBlock control_net;
+    std::string weight_prefix;
 
     ggml_backend_buffer_t control_buffer = nullptr;
     ggml_context* control_ctx            = nullptr;
@@ -321,9 +322,10 @@ struct ControlNet : public GGMLRunner {
     ControlNet(ggml_backend_t backend,
                ggml_backend_t params_backend,
                const String2TensorStorage& tensor_storage_map = {},
-               SDVersion version                              = VERSION_SD1)
-        : GGMLRunner(backend, params_backend), control_net(version) {
-        control_net.init(params_ctx, tensor_storage_map, "");
+               SDVersion version                              = VERSION_SD1,
+               const std::string& prefix                      = "")
+        : GGMLRunner(backend, params_backend), version(version), control_net(version), weight_prefix(prefix) {
+        control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
     ~ControlNet() override {
@@ -374,8 +376,8 @@ struct ControlNet : public GGMLRunner {
         return "control_net";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        control_net.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        control_net.get_param_tensors(tensors, weight_prefix);
     }
 
     ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@@ -435,7 +437,7 @@ struct ControlNet : public GGMLRunner {
             return build_graph(x, hint, timesteps, context, y);
         };
 
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
+        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false);
         if (!compute_result.has_value()) {
             return std::nullopt;
         }
@@ -472,7 +474,8 @@ struct ControlNet : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
+        model_loader.set_n_threads(n_threads);
+        bool success = model_loader.load_tensors(tensors, ignore_tensors);
 
         if (!success) {
             LOG_ERROR("load control net tensors from model loader failed");

src/model/diffusion/ernie_image.hpp

@@ -440,7 +440,7 @@ namespace ErnieImage {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/flux.hpp

@@ -1500,7 +1500,7 @@ namespace Flux {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
             };
 
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
             return result;
         }
 

src/model/diffusion/hidream_o1.hpp

@@ -323,11 +323,15 @@ namespace HiDreamO1 {
             return gf;
         }
 
-        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& image,
+                                  bool auto_free           = true,
+                                  bool free_compute_buffer = true,
+                                  bool free_compute_params = true) {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
+            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -455,7 +459,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -494,21 +498,6 @@ namespace HiDreamO1 {
             vision_runner->get_param_tensors(tensors);
         }
 
-        bool alloc_params_buffer() override {
-            if (!vision_runner->alloc_params_buffer()) {
-                return false;
-            }
-            return true;
-        }
-
-        void free_params_buffer() override {
-            vision_runner->free_params_buffer();
-        }
-
-        size_t get_params_buffer_size() override {
-            return vision_runner->get_params_buffer_size();
-        }
-
         void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
             vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
         }
@@ -521,6 +510,14 @@ namespace HiDreamO1 {
             vision_runner->set_weight_adapter(adapter);
         }
 
+        void set_weight_manager(const std::shared_ptr<RunnerWeightManager>& manager) override {
+            vision_runner->set_weight_manager(manager);
+        }
+
+        void runner_done() override {
+            vision_runner->runner_done();
+        }
+
         SDCondition get_learned_condition(int n_threads,
                                           const ConditionerParams& conditioner_params) override {
             SDCondition result;
@@ -666,7 +663,7 @@ namespace HiDreamO1 {
             result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
             result.c_image_embeds.reserve(vlm_images.size());
             for (const auto& vlm_image : vlm_images) {
-                auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true);
                 if (image_embed.empty()) {
                     LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
                     return SDCondition();

src/model/diffusion/ideogram4.hpp

@@ -537,7 +537,7 @@ namespace Ideogram4 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, use_uncond_model);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lens.hpp

@@ -408,7 +408,7 @@ namespace Lens {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/ltxv.hpp

@@ -1939,7 +1939,7 @@ namespace LTXV {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
             };
-            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
             return out;
         }
 

src/model/diffusion/mmdit.hpp

@@ -935,7 +935,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/pid.hpp

@@ -823,7 +823,7 @@ namespace Pid {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/qwen_image.hpp

@@ -627,7 +627,7 @@ namespace Qwen {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/unet.hpp

@@ -772,7 +772,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/wan.hpp

@@ -950,7 +950,7 @@ namespace WAN {
                 return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/z_image.hpp

@@ -634,7 +634,7 @@ namespace ZImage {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,