use actual scheduler sigmas for ucache bounds

rmatif · rmatif · commit 7c3dd7e6ca86 · 2025-12-10T20:07:04.000Z
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -1558,6 +1558,10 @@ class StableDiffusionGGML {
             }
         }
 
+        if (ucache_enabled) {
+            ucache_state.set_sigmas(sigmas);
+        }
+
         size_t steps          = sigmas.size() - 1;
         struct ggml_tensor* x = ggml_dup_tensor(work_ctx, init_latent);
         copy_ggml_tensor(x, init_latent);
diff --git a/ucache.hpp b/ucache.hpp
@@ -126,6 +126,26 @@ struct UCacheState {
         }
     }
 
+    void set_sigmas(const std::vector<float>& sigmas) {
+        if (!initialized || sigmas.size() < 2) {
+            return;
+        }
+        size_t n_steps = sigmas.size() - 1;
+
+        size_t start_step = static_cast<size_t>(config.start_percent * n_steps);
+        size_t end_step   = static_cast<size_t>(config.end_percent * n_steps);
+
+        if (start_step >= n_steps) start_step = n_steps - 1;
+        if (end_step >= n_steps) end_step = n_steps - 1;
+
+        start_sigma = sigmas[start_step];
+        end_sigma   = sigmas[end_step];
+
+        if (start_sigma < end_sigma) {
+            std::swap(start_sigma, end_sigma);
+        }
+    }
+
     bool enabled() const {
         return initialized && config.enabled;
     }

Original file line number	Diff line number	Diff line change
`@@ -1558,6 +1558,10 @@ class StableDiffusionGGML {`
`1558`	`1558`	`}`
`1559`	`1559`	`}`
`1560`	`1560`
	`1561`	`+ if (ucache_enabled) {`
	`1562`	`+ ucache_state.set_sigmas(sigmas);`
	`1563`	`+ }`
	`1564`	`+`
`1561`	`1565`	`size_t steps = sigmas.size() - 1;`
`1562`	`1566`	`struct ggml_tensor* x = ggml_dup_tensor(work_ctx, init_latent);`
`1563`	`1567`	`copy_ggml_tensor(x, init_latent);`