refactor: unify runner lifecycles and weight residency (#1940)

src/model/adapter/ip_adapter.hpp

@@ -200,7 +200,7 @@ namespace IPAdapter {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(image_embeds);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
         }
     };
 

src/model/adapter/lora.hpp

@@ -125,13 +125,12 @@ struct LoraModel : public GGMLRunner {
     }
 
     void release_loaded_tensors() {
-        runner_done();
-        free_compute_buffer();
+        runner_end();
         model_manager.reset();
         free_params_ctx();
         alloc_params_ctx();
-        model_manager  = std::make_shared<ModelManager>();
-        weight_manager = model_manager;
+        model_manager     = std::make_shared<ModelManager>();
+        residency_manager = model_manager;
         lora_tensors.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
@@ -952,16 +951,19 @@ struct LoraModel : public GGMLRunner {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_lora_graph(model_tensors, model_tensor_names, version);
         };
-        GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
-        stat(!warn_unused);
-        for (auto item : original_tensor_to_final_tensor) {
-            ggml_tensor* original_tensor = item.first;
-            ggml_tensor* final_tensor    = item.second;
-
-            ggml_backend_tensor_copy(final_tensor, original_tensor);
+        auto read_outputs = [&]() {
+            for (const auto& item : original_tensor_to_final_tensor) {
+                ggml_backend_tensor_copy(item.second, item.first);
+            }
+            return true;
+        };
+        auto result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        if (!result.has_value()) {
+            LOG_ERROR("LoRA graph execution failed");
         }
+        stat(!warn_unused);
         original_tensor_to_final_tensor.clear();
-        GGMLRunner::free_compute_buffer();
+        runner_end();
     }
 
     void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {

src/model/adapter/pmid.hpp

@@ -558,7 +558,7 @@ public:
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
     }
 };
 

src/model/detector/yolov8.h

@@ -355,7 +355,7 @@ struct YOLOv8Runner : public GGMLRunner {
 
     sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
         auto get_graph = [&]() { return build_graph(input); };
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
     }
 };
 

src/model/diffusion/anima.hpp

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

src/model/diffusion/boogu.hpp

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

src/model/diffusion/control.hpp

@@ -423,19 +423,26 @@ struct ControlNet : public GGMLRunner {
             return build_graph(x, hint, timesteps, context, y);
         };
 
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
+        auto read_outputs = [&]() {
+            controls.clear();
+            controls.reserve(control_outputs_ggml.size());
+            for (ggml_tensor* control : control_outputs_ggml) {
+                auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
+                if (control_host.empty()) {
+                    return false;
+                }
+                controls.push_back(std::move(control_host));
+            }
+            return true;
+        };
+        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, true, read_outputs);
+        control_outputs_ggml.clear();
+        guided_hint_output_ggml = nullptr;
         if (!compute_result.has_value()) {
+            controls.clear();
             return std::nullopt;
         }
-
         guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
-        controls.clear();
-        controls.reserve(control_outputs_ggml.size());
-        for (ggml_tensor* control : control_outputs_ggml) {
-            auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
-            GGML_ASSERT(!control_host.empty());
-            controls.push_back(std::move(control_host));
-        }
         return controls;
     }
 
@@ -444,10 +451,10 @@ struct ControlNet : public GGMLRunner {
         std::map<std::string, ggml_tensor*> tensors;
         control_net.get_param_tensors(tensors);
 
-        auto manager = std::dynamic_pointer_cast<ModelManager>(weight_manager.lock());
+        auto manager = std::dynamic_pointer_cast<ModelManager>(residency_manager.lock());
         if (manager == nullptr) {
             owned_model_manager = std::make_shared<ModelManager>();
-            weight_manager      = owned_model_manager;
+            residency_manager   = owned_model_manager;
             manager             = owned_model_manager;
         }
 

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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/flux.hpp

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

src/model/diffusion/hidream_o1.hpp

@@ -325,13 +325,11 @@ namespace HiDreamO1 {
 
         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) {
+                                  bool auto_runner_end = true) {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
+            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
@@ -459,7 +457,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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -510,8 +508,8 @@ namespace HiDreamO1 {
             vision_runner->set_weight_adapter(adapter);
         }
 
-        void runner_done() override {
-            vision_runner->runner_done();
+        void runner_end() override {
+            vision_runner->runner_end();
         }
 
         SDCondition get_learned_condition(int n_threads,
@@ -659,7 +657,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, false, true, true);
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false);
                 if (image_embed.empty()) {
                     LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
                     return SDCondition();

src/model/diffusion/hunyuan.hpp

@@ -654,7 +654,7 @@ namespace Hunyuan {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/krea2.hpp

@@ -775,7 +775,7 @@ namespace Krea2 {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, ref_latents, ref_image_params);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/lingbot_video.hpp

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

src/model/diffusion/ltxv.hpp

@@ -1998,7 +1998,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, false, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
             return out;
         }
 

src/model/diffusion/mage_flow.hpp

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

src/model/diffusion/minimax_h3.hpp

@@ -1168,8 +1168,6 @@ namespace MiniMaxH3 {
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
                                                                               n_threads,
-                                                                              false,
-                                                                              false,
                                                                               false),
                                                    params.x->dim());
         }

src/model/diffusion/minit2i.hpp

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

src/model/diffusion/mmdit.hpp

@@ -987,7 +987,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, false, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,

src/model/diffusion/pid.hpp

@@ -938,7 +938,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, false, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,

src/model/diffusion/qwen_image.hpp

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

src/model/diffusion/unet.hpp

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

src/model/diffusion/z_image.hpp

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

src/model/te/clip.hpp

@@ -568,13 +568,11 @@ struct CLIPTextModelRunner : public GGMLRunner {
                               size_t max_token_idx,
                               bool return_pooled,
                               int clip_skip,
-                              bool auto_free           = true,
-                              bool free_compute_buffer = true,
-                              bool free_compute_params = true) {
+                              bool auto_runner_end = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
+        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }

src/model/te/llm.hpp

@@ -2079,9 +2079,7 @@ namespace LLM {
                                   const ImageEmbeds& image_embeds,
                                   std::set<int> out_layers,
                                   bool return_all_hidden_states                      = false,
-                                  bool auto_free                                     = true,
-                                  bool free_compute_buffer                           = true,
-                                  bool free_compute_params                           = true,
+                                  bool auto_runner_end                               = true,
                                   const DeepStackImageEmbeds& deepstack_image_embeds = {},
                                   const std::vector<ImageGrid>& image_grids          = {}) {
             auto get_graph = [&]() -> ggml_cgraph* {
@@ -2093,7 +2091,7 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params),
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end),
                                                    input_ids.dim() + 1);
         }
 
@@ -2173,13 +2171,11 @@ namespace LLM {
 
         sd::Tensor<float> encode_image(const int n_threads,
                                        const sd::Tensor<float>& image,
-                                       bool auto_free           = false,
-                                       bool free_compute_buffer = false,
-                                       bool free_compute_params = false) {
+                                       bool auto_runner_end = false) {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
         }
 
         ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
@@ -2287,13 +2283,11 @@ namespace LLM {
 
         std::vector<sd::Tensor<float>> encode_image_outputs(const int n_threads,
                                                             const sd::Tensor<float>& image,
-                                                            bool auto_free           = false,
-                                                            bool free_compute_buffer = false,
-                                                            bool free_compute_params = false) {
+                                                            bool auto_runner_end = false) {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_outputs_graph(image);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
+            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }
@@ -2312,20 +2306,14 @@ namespace LLM {
 
         std::vector<sd::Tensor<float>> encode_video_block_outputs(const int n_threads,
                                                                   const sd::Tensor<float>& frames,
-                                                                  bool auto_free           = false,
-                                                                  bool free_compute_buffer = false,
-                                                                  bool free_compute_params = false) {
+                                                                  bool auto_runner_end = false) {
             int grid_h        = static_cast<int>(frames.shape()[1] / config.vision.patch_size);
             int grid_w        = static_cast<int>(frames.shape()[0] / config.vision.patch_size);
             auto pixel_values = process_video_block_tensor(frames, config.vision);
             auto get_graph    = [&]() -> ggml_cgraph* {
                 return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
             };
-            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph,
-                                                                     n_threads,
-                                                                     auto_free,
-                                                                     free_compute_buffer,
-                                                                     free_compute_params));
+            auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end));
             if (combined.empty()) {
                 return {};
             }

src/model/te/t5.hpp

@@ -451,13 +451,11 @@ struct T5Runner : public GGMLRunner {
     sd::Tensor<float> compute(const int n_threads,
                               const sd::Tensor<int32_t>& input_ids,
                               const sd::Tensor<float>& attention_mask,
-                              bool auto_free           = true,
-                              bool free_compute_buffer = true,
-                              bool free_compute_params = true) {
+                              bool auto_runner_end = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_runner_end), 3);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,

src/model/upscaler/esrgan.hpp

@@ -265,7 +265,7 @@ struct ESRGAN : public GGMLRunner {
     sd::Tensor<float> compute(const int n_threads,
                               const sd::Tensor<float>& x) {
         auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
-        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         return result;
     }
 };

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -499,7 +499,7 @@ namespace LTXVUpsampler {
             }
             size_t expected_dim = static_cast<size_t>(x.dim());
             auto get_graph      = [&]() -> ggml_cgraph* { return build_graph(x); };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
         }
     };
 

src/model/vae/auto_encoder_kl.hpp

@@ -744,7 +744,7 @@ struct AutoEncoderKL : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
     }
 
     sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {

src/model/vae/hunyuan_vae.hpp

@@ -827,9 +827,7 @@ namespace Hunyuan {
             };
             auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
                                                                                      n_threads,
-                                                                                     true,
-                                                                                     true,
-                                                                                     true),
+                                                                                     false),
                                                           graph_input.dim());
             if (!output.empty() && input.dim() == 4) {
                 output.squeeze_(2);

src/model/vae/ltx_audio_vae.hpp

@@ -1042,7 +1042,7 @@ namespace LTXV {
                 ggml_build_forward_expand(gf, waveform);
                 return gf;
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), 4);
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
             int64_t t1  = ggml_time_ms();
             LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
             return result;

src/model/vae/ltx_vae.hpp

@@ -1334,7 +1334,6 @@ struct LTXVideoVAE : public VAE {
                   (int)plan.tiles.size());
 
         free_cache_ctx_and_buffer();
-        cache_tensor_map.clear();
 
         auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& z_chunk, const VAETemporalTile& tile) {
             LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
@@ -1349,12 +1348,11 @@ struct LTXVideoVAE : public VAE {
                                                  static_cast<int>(tile.start),
                                                  tile.overlap);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
                                                    expected_dim);
         });
 
         free_cache_ctx_and_buffer();
-        cache_tensor_map.clear();
         return output;
     }
 
@@ -1407,7 +1405,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input, decode_graph);
         };
-        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
+        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
         if (result.empty()) {
             return {};
         }
@@ -1420,7 +1418,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_latent_statistics_graph(z, normalize);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
                                                static_cast<size_t>(z.dim()));
     }
 

src/model/vae/mage_vae.hpp

@@ -490,7 +490,7 @@ namespace MageVAE {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input, decode_graph);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), input.dim());
         }
 
         int get_encoder_output_channels(int input_channels) override {

src/model/vae/minimax_h3_audio_vae.hpp

@@ -480,7 +480,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
+                GGMLRunner::compute<float>(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
@@ -500,7 +500,7 @@ namespace MiniMaxH3 {
                 return graph;
             };
             auto result = restore_trailing_singleton_dims(
-                GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
+                GGMLRunner::compute<float>(get_graph, n_threads, false),
                 4);
             int64_t t1 = ggml_time_ms();
             LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",

src/model/vae/minimax_h3_vae.hpp

@@ -793,8 +793,6 @@ namespace MiniMaxH3VAE {
             return restore_trailing_singleton_dims(
                 GGMLRunner::compute<float>(get_graph,
                                            n_threads,
-                                           false,
-                                           false,
                                            false),
                 5);
         }

src/model/vae/tae.hpp

@@ -787,7 +787,7 @@ struct TinyImageAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 
@@ -872,7 +872,7 @@ struct TinyVideoAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
     }
 };
 

src/model/vae/vae.hpp

@@ -251,7 +251,7 @@ public:
                                                   tiling_params);
         }
 
-        runner_done();
+        runner_end();
 
         if (output.empty()) {
             LOG_ERROR("vae encode compute failed");
@@ -305,7 +305,7 @@ public:
                                                   tiling_params);
         }
 
-        runner_done();
+        runner_end();
 
         if (output.empty()) {
             LOG_ERROR("vae decode compute failed");

src/model/vae/wan_vae.hpp

@@ -1415,7 +1415,6 @@ namespace WAN {
                       (int)plan.tiles.size());
 
             free_cache_ctx_and_buffer();
-            cache_tensor_map.clear();
             ae.clear_cache();
 
             auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor<float>& input_tile, const VAETemporalTile& tile) {
@@ -1428,12 +1427,11 @@ namespace WAN {
                     return build_temporal_tile_graph(input_tile, static_cast<int>(tile.start));
                 };
                 return restore_trailing_singleton_dims(
-                    GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
+                    GGMLRunner::compute<float>(get_graph, n_threads, false),
                     static_cast<size_t>(input.dim()));
             });
 
             free_cache_ctx_and_buffer();
-            cache_tensor_map.clear();
             ae.clear_cache();
             return output;
         }
@@ -1448,7 +1446,7 @@ namespace WAN {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input.empty() ? z : input, decode_graph);
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
                                                           input.empty() ? z.dim() : input.dim());
             if (!result.empty() && z.dim() == 4) {
                 result.squeeze_(2);