refactor: route all runner params through model manager (#1649)

src/model/vae/auto_encoder_kl.hpp

@@ -213,9 +213,9 @@ protected:
         params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
     }
 
-    float get_alpha() {
-        float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]);
-        return sigmoid(alpha);
+    ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
+        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 
 public:
@@ -250,10 +250,12 @@ public:
 
         x = time_stack->forward(ctx, x);  // b t c (h w)
 
-        float alpha = get_alpha();
-        x           = ggml_add(ctx->ggml_ctx,
-                               ggml_ext_scale(ctx->ggml_ctx, x, alpha),
-                               ggml_ext_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha));
+        auto alpha = get_alpha(ctx);
+        x          = ggml_add(ctx->ggml_ctx,
+                              x_mix,
+                              ggml_mul(ctx->ggml_ctx,
+                                       ggml_sub(ctx->ggml_ctx, x, x_mix),
+                                       alpha));
 
         x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // b c t (h w) -> b t c (h w)
         x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, C, T * B);                     // b t c (h w) -> (b t) c h w
@@ -664,13 +666,13 @@ struct AutoEncoderKL : public VAE {
     AutoEncoderKLModel ae;
 
     AutoEncoderKL(ggml_backend_t backend,
-                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map,
                   const std::string prefix,
-                  bool decode_only       = false,
-                  bool use_video_decoder = false,
-                  SDVersion version      = VERSION_SD1)
-        : VAE(version, backend, params_backend, prefix), decode_only(decode_only) {
+                  bool decode_only                                    = false,
+                  bool use_video_decoder                              = false,
+                  SDVersion version                                   = VERSION_SD1,
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) {
         if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
             scale_factor = 0.18215f;
             shift_factor = 0.f;

src/model/vae/ltx_audio_vae.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
+#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 #define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 
 #include <cmath>
@@ -9,6 +9,7 @@
 
 #include "core/ggml_extend.hpp"
 #include "model_loader.h"
+#include "model_manager.h"
 
 namespace LTXV {
 
@@ -1001,10 +1002,10 @@ namespace LTXV {
         sd::Tensor<float> bwe_skip_filter_tensor;
 
         LTXAudioVAERunner(ggml_backend_t backend,
-                          ggml_backend_t params_backend,
                           const String2TensorStorage& tensor_storage_map,
-                          const std::string& prefix = "")
-            : GGMLRunner(backend, params_backend),
+                          const std::string& prefix                           = "",
+                          std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
               weight_prefix(prefix),
               config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
               model(config) {
@@ -1019,7 +1020,7 @@ namespace LTXV {
             model.get_param_tensors(tensors, weight_prefix);
         }
 
-        size_t get_params_buffer_size() {
+        size_t get_params_mem_size() {
             return model.get_params_mem_size();
         }
 
@@ -1066,7 +1067,8 @@ namespace LTXV {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file(model_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -1074,20 +1076,17 @@ namespace LTXV {
 
             auto& tensor_storage_map = model_loader.get_tensor_storage_map();
             auto ltx_audio_vae       = std::make_shared<LTXAudioVAERunner>(backend,
-                                                                     backend,
                                                                      tensor_storage_map,
-                                                                     prefix);
+                                                                     prefix,
+                                                                     model_manager);
 
-            if (!ltx_audio_vae->alloc_params_buffer()) {
-                LOG_ERROR("ltx audio vae buffer allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            ltx_audio_vae->get_param_tensors(tensors);
-
-            if (!model_loader.load_tensors(tensors)) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("LTX audio VAE test",
+                                                       *ltx_audio_vae,
+                                                       ModelManager::ResidencyMode::Resident,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register ltx audio vae tensors with model manager failed");
                 return;
             }
 

src/model/vae/ltx_vae.hpp

@@ -957,8 +957,8 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning) {
-                auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]);
-                scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier);
+                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
             x = conv_in->forward(ctx, x, causal_decoder);
@@ -1008,8 +1008,8 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning && timestep != nullptr) {
-                auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]);
-                scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier);
+                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
             // conv_in with feat_map for left temporal context
@@ -1223,11 +1223,11 @@ struct LTXVideoVAE : public VAE {
     LTXVAE::VideoVAE vae;
 
     LTXVideoVAE(ggml_backend_t backend,
-                ggml_backend_t params_backend,
                 const String2TensorStorage& tensor_storage_map,
                 const std::string& prefix,
-                bool decode_only  = true,
-                SDVersion version = VERSION_LTXAV)
+                bool decode_only                                    = true,
+                SDVersion version                                   = VERSION_LTXAV,
+                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
         : decode_only(decode_only),
           ltx_vae_version(LTXVAE::detect_ltx_vae_version(tensor_storage_map, prefix)),
           timestep_conditioning(LTXVAE::detect_ltx_vae_timestep_conditioning(tensor_storage_map, prefix)),
@@ -1239,7 +1239,7 @@ struct LTXVideoVAE : public VAE {
               patch_size,
               tensor_storage_map,
               prefix),
-          VAE(version, backend, params_backend, prefix) {
+          VAE(version, backend, prefix, weight_manager) {
         vae.init(params_ctx, tensor_storage_map, prefix);
         decode_timestep_tensor.values()[0] = vae.decode_timestep;
     }
@@ -1521,7 +1521,8 @@ struct LTXVideoVAE : public VAE {
         ggml_backend_t backend = sd_backend_cpu_init();
         LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
 
-        ModelLoader model_loader;
+        auto model_manager        = std::make_shared<ModelManager>();
+        ModelLoader& model_loader = model_manager->loader();
         if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) {
             LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
             return;
@@ -1529,22 +1530,19 @@ struct LTXVideoVAE : public VAE {
 
         auto& tensor_storage_map         = model_loader.get_tensor_storage_map();
         std::shared_ptr<LTXVideoVAE> vae = std::make_shared<LTXVideoVAE>(backend,
-                                                                         backend,
                                                                          tensor_storage_map,
                                                                          "first_stage_model",
                                                                          true,
-                                                                         VERSION_LTXAV);
+                                                                         VERSION_LTXAV,
+                                                                         model_manager);
 
-        if (!vae->alloc_params_buffer()) {
-            LOG_ERROR("vae buffer allocation failed");
-            return;
-        }
-
-        std::map<std::string, ggml_tensor*> tensors;
-        vae->get_param_tensors(tensors);
-
-        if (!model_loader.load_tensors(tensors)) {
-            LOG_ERROR("load tensors from model loader failed");
+        if (!model_manager->register_runner_params("LTX VAE test",
+                                                   *vae,
+                                                   ModelManager::ResidencyMode::Resident,
+                                                   backend,
+                                                   backend) ||
+            !model_manager->validate_registered_tensors()) {
+            LOG_ERROR("register ltx vae tensors with model manager failed");
             return;
         }
 

src/model/vae/tae.hpp

@@ -623,12 +623,12 @@ struct TinyImageAutoEncoder : public VAE {
     bool decode_only = false;
 
     TinyImageAutoEncoder(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_SD1)
-        : VAE(version, backend, params_backend, "tae"),
+                         bool decoder_only                                   = true,
+                         SDVersion version                                   = VERSION_SD1,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "tae", weight_manager),
           decode_only(decoder_only),
           taesd(decoder_only, version) {
         scale_input = false;
@@ -686,12 +686,12 @@ struct TinyVideoAutoEncoder : public VAE {
     bool is_wide     = false;
 
     TinyVideoAutoEncoder(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_WAN2)
-        : VAE(version, backend, params_backend, "tae"),
+                         bool decoder_only                                   = true,
+                         SDVersion version                                   = VERSION_WAN2,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "tae", weight_manager),
           decode_only(decoder_only) {
         for (auto tensor_storage : tensor_storage_map) {
             if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {

src/model/vae/vae.hpp

@@ -1,8 +1,9 @@
-#ifndef __SD_MODEL_VAE_VAE_HPP__
+#ifndef __SD_MODEL_VAE_VAE_HPP__
 #define __SD_MODEL_VAE_VAE_HPP__
 
 #include "core/tensor_ggml.hpp"
 #include "model/common/block.hpp"
+#include "model_manager.h"
 
 struct VAE : public GGMLRunner {
 protected:
@@ -63,8 +64,11 @@ protected:
     }
 
 public:
-    VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend, const std::string& weight_prefix = "")
-        : version(version), weight_prefix(weight_prefix), GGMLRunner(backend, params_backend) {}
+    VAE(SDVersion version,
+        ggml_backend_t backend,
+        const std::string& weight_prefix                    = "",
+        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : version(version), weight_prefix(weight_prefix), GGMLRunner(backend, weight_manager) {}
 
     int get_scale_factor() {
         int scale_factor = 8;
@@ -224,8 +228,10 @@ public:
 };
 
 struct FakeVAE : public VAE {
-    FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
-        : VAE(version, backend, params_backend) {}
+    FakeVAE(SDVersion version,
+            ggml_backend_t backend,
+            std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "", weight_manager) {}
 
     int get_encoder_output_channels(int input_channels) {
         return input_channels;

src/model/vae/wan_vae.hpp

@@ -1124,12 +1124,12 @@ namespace WAN {
         WanVAE ae;
 
         WanVAERunner(ggml_backend_t backend,
-                     ggml_backend_t params_backend,
-                     const String2TensorStorage& tensor_storage_map = {},
-                     const std::string prefix                       = "",
-                     bool decode_only                               = false,
-                     SDVersion version                              = VERSION_WAN2)
-            : VAE(version, backend, params_backend, prefix), decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V) {
+                     const String2TensorStorage& tensor_storage_map      = {},
+                     const std::string prefix                            = "",
+                     bool decode_only                                    = false,
+                     SDVersion version                                   = VERSION_WAN2,
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : VAE(version, backend, prefix, weight_manager), decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V) {
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1327,27 +1327,24 @@ namespace WAN {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             ggml_backend_t backend            = sd_backend_cpu_init();
             ggml_type model_data_type         = GGML_TYPE_F16;
-            std::shared_ptr<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, backend, String2TensorStorage{}, "first_stage_model", false, VERSION_WAN2_2_TI2V);
+            auto model_manager                = std::make_shared<ModelManager>();
+            std::shared_ptr<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, String2TensorStorage{}, "first_stage_model", false, VERSION_WAN2_2_TI2V, model_manager);
             {
                 LOG_INFO("loading from '%s'", file_path.c_str());
 
-                if (!vae->alloc_params_buffer()) {
-                    LOG_ERROR("vae buffer allocation failed");
-                    return;
-                }
-                std::map<std::string, ggml_tensor*> tensors;
-                vae->get_param_tensors(tensors);
-
-                ModelLoader model_loader;
+                ModelLoader& model_loader = model_manager->loader();
                 if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) {
                     LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                     return;
                 }
 
-                bool success = model_loader.load_tensors(tensors);
-
-                if (!success) {
-                    LOG_ERROR("load tensors from model loader failed");
+                if (!model_manager->register_runner_params("Wan VAE test",
+                                                           *vae,
+                                                           ModelManager::ResidencyMode::Resident,
+                                                           backend,
+                                                           backend) ||
+                    !model_manager->validate_registered_tensors()) {
+                    LOG_ERROR("register wan vae tensors with model manager failed");
                     return;
                 }