add decay rate and relative threshold

rmatif · rmatif · commit 86d176b531d4 · 2025-12-10T11:06:44.000Z
diff --git a/examples/cli/README.md b/examples/cli/README.md
@@ -125,5 +125,7 @@ Generation Options:
   --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
   -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
   --cache-mode                             caching method: 'easycache' for DiT models, 'ucache' for UNET models (SD1.x/SD2.x/SDXL)
-  --cache-option                           cache parameters "threshold,start_percent,end_percent" (default: 0.2,0.15,0.95 for easycache, 1.0,0.15,0.95 for ucache)
+  --cache-option                           cache parameters: easycache uses "threshold,start,end" (default: 0.2,0.15,0.95).
+                                           ucache uses "threshold,start,end[,decay,relative]" (default: 1.0,0.15,0.95,1.0,1).
+                                           decay: error decay rate (0.0-1.0), relative: use relative threshold (0 or 1)
 ```
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -1441,7 +1441,7 @@ struct SDGenerationParams {
              on_cache_mode_arg},
             {"",
              "--cache-option",
-             "cache parameters \"threshold,start_percent,end_percent\" (default: 0.2,0.15,0.95 for easycache, 1.0,0.15,0.95 for ucache)",
+             "cache parameters \"threshold,start,end[,warmup,decay,relative]\" (ucache extended: warmup=0, decay=1.0, relative=1)",
              on_cache_option_arg},
 
         };
@@ -1568,28 +1568,32 @@ struct SDGenerationParams {
                 }
             }
 
-            float values[3] = {0.0f, 0.0f, 0.0f};
+            // Format: threshold,start,end[,decay,relative]
+            // - values[0-2]: threshold, start_percent, end_percent (required)
+            // - values[3]: error_decay_rate (optional, default: 1.0)
+            // - values[4]: use_relative_threshold (optional, 0 or 1, default: 1)
+            float values[5] = {0.0f, 0.0f, 0.0f, 1.0f, 1.0f};
             std::stringstream ss(option_str);
             std::string token;
             int idx = 0;
+            auto trim = [](std::string& s) {
+                const char* whitespace = " \t\r\n";
+                auto start             = s.find_first_not_of(whitespace);
+                if (start == std::string::npos) {
+                    s.clear();
+                    return;
+                }
+                auto end = s.find_last_not_of(whitespace);
+                s        = s.substr(start, end - start + 1);
+            };
             while (std::getline(ss, token, ',')) {
-                auto trim = [](std::string& s) {
-                    const char* whitespace = " \t\r\n";
-                    auto start             = s.find_first_not_of(whitespace);
-                    if (start == std::string::npos) {
-                        s.clear();
-                        return;
-                    }
-                    auto end = s.find_last_not_of(whitespace);
-                    s        = s.substr(start, end - start + 1);
-                };
                 trim(token);
                 if (token.empty()) {
                     fprintf(stderr, "error: invalid cache option '%s'\n", option_str.c_str());
                     return false;
                 }
-                if (idx >= 3) {
-                    fprintf(stderr, "error: cache option expects exactly 3 comma-separated values (threshold,start,end)\n");
+                if (idx >= 5) {
+                    fprintf(stderr, "error: cache option expects 3-5 comma-separated values (threshold,start,end[,decay,relative])\n");
                     return false;
                 }
                 try {
@@ -1600,8 +1604,8 @@ struct SDGenerationParams {
                 }
                 idx++;
             }
-            if (idx != 3) {
-                fprintf(stderr, "error: cache option expects exactly 3 comma-separated values (threshold,start,end)\n");
+            if (idx < 3) {
+                fprintf(stderr, "error: cache option expects at least 3 comma-separated values (threshold,start,end)\n");
                 return false;
             }
             if (values[0] < 0.0f) {
@@ -1619,10 +1623,12 @@ struct SDGenerationParams {
                 easycache_params.start_percent   = values[1];
                 easycache_params.end_percent     = values[2];
             } else {
-                ucache_params.enabled         = true;
-                ucache_params.reuse_threshold = values[0];
-                ucache_params.start_percent   = values[1];
-                ucache_params.end_percent     = values[2];
+                ucache_params.enabled                = true;
+                ucache_params.reuse_threshold        = values[0];
+                ucache_params.start_percent          = values[1];
+                ucache_params.end_percent            = values[2];
+                ucache_params.error_decay_rate       = values[3];
+                ucache_params.use_relative_threshold = (values[4] != 0.0f);
             }
         }
 
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -1538,10 +1538,12 @@ class StableDiffusionGGML {
                 LOG_WARN("UCache requested but not supported for this model type (only UNET models)");
             } else {
                 UCacheConfig ucache_config;
-                ucache_config.enabled         = true;
-                ucache_config.reuse_threshold = std::max(0.0f, ucache_params->reuse_threshold);
-                ucache_config.start_percent   = ucache_params->start_percent;
-                ucache_config.end_percent     = ucache_params->end_percent;
+                ucache_config.enabled                = true;
+                ucache_config.reuse_threshold        = std::max(0.0f, ucache_params->reuse_threshold);
+                ucache_config.start_percent          = ucache_params->start_percent;
+                ucache_config.end_percent            = ucache_params->end_percent;
+                ucache_config.error_decay_rate       = std::max(0.0f, std::min(1.0f, ucache_params->error_decay_rate));
+                ucache_config.use_relative_threshold = ucache_params->use_relative_threshold;
                 bool percent_valid            = ucache_config.start_percent >= 0.0f &&
                                      ucache_config.start_percent < 1.0f &&
                                      ucache_config.end_percent > 0.0f &&
@@ -1555,10 +1557,12 @@ class StableDiffusionGGML {
                     ucache_state.init(ucache_config, denoiser.get());
                     if (ucache_state.enabled()) {
                         ucache_enabled = true;
-                        LOG_INFO("UCache enabled - threshold: %.3f, start_percent: %.2f, end_percent: %.2f",
+                        LOG_INFO("UCache enabled - threshold: %.3f, start: %.2f, end: %.2f, decay: %.2f, relative: %s",
                                  ucache_config.reuse_threshold,
                                  ucache_config.start_percent,
-                                 ucache_config.end_percent);
+                                 ucache_config.end_percent,
+                                 ucache_config.error_decay_rate,
+                                 ucache_config.use_relative_threshold ? "true" : "false");
                     } else {
                         LOG_WARN("UCache requested but could not be initialized for this run");
                     }
@@ -2594,11 +2598,13 @@ void sd_easycache_params_init(sd_easycache_params_t* easycache_params) {
 }
 
 void sd_ucache_params_init(sd_ucache_params_t* ucache_params) {
-    *ucache_params                 = {};
-    ucache_params->enabled         = false;
-    ucache_params->reuse_threshold = 1.0f;
-    ucache_params->start_percent   = 0.15f;
-    ucache_params->end_percent     = 0.95f;
+    *ucache_params                        = {};
+    ucache_params->enabled                = false;
+    ucache_params->reuse_threshold        = 1.0f;
+    ucache_params->start_percent          = 0.15f;
+    ucache_params->end_percent            = 0.95f;
+    ucache_params->error_decay_rate       = 1.0f;
+    ucache_params->use_relative_threshold = true;
 }
 
 void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -246,6 +246,8 @@ typedef struct {
     float reuse_threshold;
     float start_percent;
     float end_percent;
+    float error_decay_rate;
+    bool use_relative_threshold;
 } sd_ucache_params_t;
 
 typedef struct {
diff --git a/ucache.hpp b/ucache.hpp
@@ -10,10 +10,15 @@
 #include "ggml_extend.hpp"
 
 struct UCacheConfig {
-    bool enabled          = false;
-    float reuse_threshold = 1.0f;
-    float start_percent   = 0.15f;
-    float end_percent     = 0.95f;
+    bool enabled                 = false;
+    float reuse_threshold        = 1.0f;
+    float start_percent          = 0.15f;
+    float end_percent            = 0.95f;
+    float error_decay_rate       = 1.0f;
+    bool use_relative_threshold  = true;
+    bool adaptive_threshold      = true;
+    float early_step_multiplier  = 0.5f;
+    float late_step_multiplier   = 1.5f;
 };
 
 struct UCacheCacheEntry {
@@ -44,6 +49,45 @@ struct UCacheState {
     bool has_last_input_change            = false;
     int total_steps_skipped               = 0;
     int current_step_index                = -1;
+    int steps_computed_since_active       = 0;
+    float accumulated_error               = 0.0f;
+    float reference_output_norm           = 0.0f;
+
+    struct BlockMetrics {
+        float sum_transformation_rate = 0.0f;
+        float sum_output_norm         = 0.0f;
+        int sample_count              = 0;
+        float min_change_rate         = std::numeric_limits<float>::max();
+        float max_change_rate         = 0.0f;
+
+        void reset() {
+            sum_transformation_rate = 0.0f;
+            sum_output_norm         = 0.0f;
+            sample_count            = 0;
+            min_change_rate         = std::numeric_limits<float>::max();
+            max_change_rate         = 0.0f;
+        }
+
+        void record(float change_rate, float output_norm) {
+            if (std::isfinite(change_rate) && change_rate > 0.0f) {
+                sum_transformation_rate += change_rate;
+                sum_output_norm += output_norm;
+                sample_count++;
+                if (change_rate < min_change_rate) min_change_rate = change_rate;
+                if (change_rate > max_change_rate) max_change_rate = change_rate;
+            }
+        }
+
+        float avg_transformation_rate() const {
+            return (sample_count > 0) ? (sum_transformation_rate / sample_count) : 0.0f;
+        }
+
+        float avg_output_norm() const {
+            return (sample_count > 0) ? (sum_output_norm / sample_count) : 0.0f;
+        }
+    };
+    BlockMetrics block_metrics;
+    int total_active_steps = 0;
 
     void reset_runtime() {
         initial_step      = true;
@@ -64,6 +108,11 @@ struct UCacheState {
         has_last_input_change            = false;
         total_steps_skipped              = 0;
         current_step_index               = -1;
+        steps_computed_since_active      = 0;
+        accumulated_error                = 0.0f;
+        reference_output_norm            = 0.0f;
+        block_metrics.reset();
+        total_active_steps = 0;
     }
 
     void init(const UCacheConfig& cfg, Denoiser* d) {
@@ -114,6 +163,7 @@ struct UCacheState {
             return;
         }
         step_active = true;
+        total_active_steps++;
     }
 
     bool step_is_active() const {
@@ -124,6 +174,31 @@ struct UCacheState {
         return enabled() && step_active && skip_current_step;
     }
 
+    float get_adaptive_threshold(int estimated_total_steps = 0) const {
+        float base_threshold = config.reuse_threshold;
+
+        if (!config.adaptive_threshold) {
+            return base_threshold;
+        }
+
+        int effective_total = estimated_total_steps;
+        if (effective_total <= 0) {
+            effective_total = std::max(20, steps_computed_since_active * 2);
+        }
+
+        float progress = (effective_total > 0) ?
+            (static_cast<float>(steps_computed_since_active) / effective_total) : 0.0f;
+
+        float multiplier = 1.0f;
+        if (progress < 0.2f) {
+            multiplier = config.early_step_multiplier;
+        } else if (progress > 0.8f) {
+            multiplier = config.late_step_multiplier;
+        }
+
+        return base_threshold * multiplier;
+    }
+
     bool has_cache(const SDCondition* cond) const {
         auto it = cache_diffs.find(cond);
         return it != cache_diffs.end() && !it->second.diff.empty();
@@ -212,15 +287,18 @@ struct UCacheState {
             last_input_change > 0.0f && output_prev_norm > 0.0f) {
 
             float approx_output_change_rate = (relative_transformation_rate * last_input_change) / output_prev_norm;
-            cumulative_change_rate += approx_output_change_rate;
+            accumulated_error = accumulated_error * config.error_decay_rate + approx_output_change_rate;
+
+            float effective_threshold = get_adaptive_threshold();
+            if (config.use_relative_threshold && reference_output_norm > 0.0f) {
+                effective_threshold = effective_threshold * reference_output_norm;
+            }
 
-            if (cumulative_change_rate < config.reuse_threshold) {
+            if (accumulated_error < effective_threshold) {
                 skip_current_step = true;
                 total_steps_skipped++;
                 apply_cache(cond, input, output);
                 return true;
-            } else {
-                cumulative_change_rate = 0.0f;
             }
         }
 
@@ -270,16 +348,31 @@ struct UCacheState {
         output_prev_norm     = (ne > 0) ? (mean_abs / static_cast<float>(ne)) : 0.0f;
         has_output_prev_norm = output_prev_norm > 0.0f;
 
+        if (reference_output_norm == 0.0f) {
+            reference_output_norm = output_prev_norm;
+        }
+
         if (has_last_input_change && last_input_change > 0.0f && output_change > 0.0f) {
             float rate = output_change / last_input_change;
             if (std::isfinite(rate)) {
                 relative_transformation_rate     = rate;
                 has_relative_transformation_rate = true;
+                block_metrics.record(rate, output_prev_norm);
             }
         }
 
-        cumulative_change_rate = 0.0f;
-        has_last_input_change  = false;
+        has_last_input_change = false;
+    }
+
+    void log_block_metrics() const {
+        if (block_metrics.sample_count > 0) {
+            LOG_INFO("UCacheBlockMetrics: samples=%d, avg_rate=%.4f, min=%.4f, max=%.4f, avg_norm=%.4f",
+                     block_metrics.sample_count,
+                     block_metrics.avg_transformation_rate(),
+                     block_metrics.min_change_rate,
+                     block_metrics.max_change_rate,
+                     block_metrics.avg_output_norm());
+        }
     }
 };
 
