feat: add configurable conditioning cache for all models (#2034)

src/pipeline/diffusion_engine.cpp

@@ -29,6 +29,7 @@
 #include "stable-diffusion.h"
 
 #include "conditioning/conditioner.hpp"
+#include "conditioning/conditioning_cache.h"
 #include "core/backend_fit.h"
 #include "extensions/generation_extension.h"
 #include "model/adapter/ip_adapter.hpp"
@@ -135,6 +136,7 @@ static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
 
 StableDiffusionGGML::StableDiffusionGGML()
     : rng(std::make_shared<PhiloxRNG>()),
+      conditioning_cache_(std::make_unique<ConditioningCache>()),
       denoiser(std::make_shared<CompVisDenoiser>()) {}
 
 StableDiffusionGGML::~StableDiffusionGGML() = default;
@@ -204,6 +206,8 @@ void StableDiffusionGGML::end_runners() {
 }
 
 bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) {
+    conditioning_cache_->clear();
+    conditioning_loras_.clear();
     end_runners();
     clear_lora_adapters();
     runtime_lora_models.clear();
@@ -915,6 +919,11 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
             return false;
         }
     }
+    if (sd_ctx_params->conditioning_cache_size < 0) {
+        LOG_ERROR("conditioning_cache_size must be non-negative");
+        return false;
+    }
+    conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
     auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
     n_threads          = sd_ctx_params->n_threads;
     tensor_executor    = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
@@ -1762,13 +1771,22 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
         extension->collect_loras(all_loras);
     }
 
-    conditioning_cache_allowed_ = all_loras.empty();
-
     int64_t t0 = ggml_time_ms();
     end_runners();
     clear_lora_adapters();
-    if (!model_manager->prepare_lora_sources(all_loras))
+    if (!model_manager->prepare_lora_sources(all_loras)) {
+        conditioning_cache_->clear();
         return false;
+    }
+    if (!std::equal(all_loras.begin(), all_loras.end(),
+                    conditioning_loras_.begin(), conditioning_loras_.end(),
+                    [](const ModelManager::LoraSpec& a, const ModelManager::LoraSpec& b) {
+                        return a.file_id == b.file_id && a.file_revision == b.file_revision &&
+                               a.multiplier == b.multiplier && a.is_high_noise == b.is_high_noise &&
+                               a.tensor_name_prefix_filter == b.tensor_name_prefix_filter;
+                    })) {
+        conditioning_cache_->clear();
+    }
     runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) {
                                   return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) {
                                       return entry.matches(spec);
@@ -1778,6 +1796,7 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
     const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras)
                                                 : apply_loras_at_runtime(all_loras);
     if (!success) {
+        conditioning_cache_->clear();
         clear_lora_adapters();
         runtime_lora_models.clear();
         return false;
@@ -1787,9 +1806,14 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
     if (!all_loras.empty()) {
         LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
     }
+    conditioning_loras_ = std::move(all_loras);
     return true;
 }
 
+SDCondition StableDiffusionGGML::get_learned_condition(const ConditionerParams& params) {
+    return conditioning_cache_->get(*cond_stage_model, n_threads, params);
+}
+
 void StableDiffusionGGML::reset_generation_extensions() {
     for (auto& extension : generation_extensions) {
         extension->reset_runtime_condition();

src/pipeline/diffusion_engine.h

@@ -26,6 +26,7 @@ class RNG;
 struct Denoiser;
 struct LoraModel;
 struct ConditionerParams;
+class ConditioningCache;
 struct SDCondition;
 struct RefImageParams;
 namespace Wav2Vec2 {
@@ -178,8 +179,9 @@ public:
     std::recursive_mutex execution_mutex;
     std::unique_ptr<ModelConfig> config_;
     RunnerState runner_state_;
-    bool conditioning_cache_allowed_ = false;
-    bool executing_                  = false;
+    std::unique_ptr<ConditioningCache> conditioning_cache_;
+    std::vector<ModelManager::LoraSpec> conditioning_loras_;
+    bool executing_ = false;
 
     std::shared_ptr<Denoiser> denoiser;
     std::vector<float> file_alphas_cumprod;
@@ -362,6 +364,8 @@ public:
 
     bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
 
+    SDCondition get_learned_condition(const ConditionerParams& params);
+
     void reset_generation_extensions();
 
     void prepare_generation_extensions(const sd_pm_params_t& pm_params,

src/pipeline/image.cpp

@@ -441,8 +441,7 @@ namespace sd::pipeline {
         sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength);
         int64_t prepare_start_ms         = ggml_time_ms();
         condition_params.zero_out_masked = false;
-        auto cond                        = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                       condition_params);
+        auto cond                        = sd->get_learned_condition(condition_params);
         if (cond.empty()) {
             LOG_ERROR("failed to encode prompt");
             return std::nullopt;
@@ -480,8 +479,7 @@ namespace sd::pipeline {
                     // LLaDA-Image CFG keeps the source latent but drops its SigVQ features.
                     condition_params.ref_images = nullptr;
                 }
-                uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                     condition_params);
+                uncond = sd->get_learned_condition(condition_params);
                 if (uncond.empty()) {
                     LOG_ERROR("failed to encode negative prompt");
                     return std::nullopt;
@@ -509,8 +507,7 @@ namespace sd::pipeline {
                 if (use_ref_latent_img_cfg) {
                     condition_params.ref_images = &empty_ref_images;
                 }
-                img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                         condition_params);
+                img_uncond = sd->get_learned_condition(condition_params);
                 if (img_uncond.empty()) {
                     LOG_ERROR("failed to encode image guidance prompt");
                     return std::nullopt;

src/pipeline/video.cpp

@@ -1157,17 +1157,12 @@ namespace sd::pipeline {
         condition_params.zero_out_masked       = true;
         condition_params.ref_images            = &latents.ref_images;
         condition_params.minimax_h3_references = &latents.minimax_presentation_refs;
-        condition_params.allow_cache =
-            sd_version_is_minimax_h3(sd->version) &&
-            sd->conditioning_cache_allowed_ &&
-            !request.use_uncond;
         if (sd_version_is_lingbot_video(sd->version) || sd_version_is_minimax_h3(sd->version)) {
             condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA;
         }
 
         int64_t prepare_start_ms = ggml_time_ms();
-        embeds.cond              = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                               condition_params);
+        embeds.cond              = sd->get_learned_condition(condition_params);
         if (embeds.cond.empty()) {
             LOG_ERROR("failed to encode video prompt");
             return std::nullopt;
@@ -1192,8 +1187,7 @@ namespace sd::pipeline {
         }
         if (request.use_uncond) {
             condition_params.text = request.negative_prompt;
-            embeds.uncond         = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                condition_params);
+            embeds.uncond         = sd->get_learned_condition(condition_params);
             if (embeds.uncond.empty()) {
                 LOG_ERROR("failed to encode negative video prompt");
                 return std::nullopt;