From e247504a6d1d1e49a9d526e153cd9ec90b894836 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20du=20Hamel?= <stephduh@live.fr>
Date: Mon, 22 Jun 2026 17:05:21 +0200
Subject: [PATCH] fix?: enforce flow shift for all schedulers

---
 src/runtime/denoiser.hpp | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp
index fed5911bc..0c73d1fd7 100644
--- a/src/runtime/denoiser.hpp
+++ b/src/runtime/denoiser.hpp
@@ -571,13 +571,19 @@ struct Denoiser {
     virtual sd::Tensor<float> inverse_noise_scaling(float sigma,
                                                     const sd::Tensor<float>& latent) = 0;
 
+    virtual float apply_shift_warp(float sigma) {
+        return sigma;
+    }
+
     virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
         auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
         std::shared_ptr<SigmaScheduler> scheduler;
+        bool shifted_scheduler = false;
         switch (scheduler_type) {
             case DISCRETE_SCHEDULER:
                 LOG_INFO("get_sigmas with discrete scheduler");
                 scheduler = std::make_shared<DiscreteScheduler>();
+                shifted_scheduler = true;
                 break;
             case KARRAS_SCHEDULER:
                 LOG_INFO("get_sigmas with Karras scheduler");
@@ -598,14 +604,17 @@ struct Denoiser {
             case SGM_UNIFORM_SCHEDULER:
                 LOG_INFO("get_sigmas with SGM Uniform scheduler");
                 scheduler = std::make_shared<SGMUniformScheduler>();
+                shifted_scheduler = true;
                 break;
             case SIMPLE_SCHEDULER:
                 LOG_INFO("get_sigmas with Simple scheduler");
                 scheduler = std::make_shared<SimpleScheduler>();
+                shifted_scheduler = true;
                 break;
             case SMOOTHSTEP_SCHEDULER:
                 LOG_INFO("get_sigmas with SmoothStep scheduler");
                 scheduler = std::make_shared<SmoothStepScheduler>();
+                shifted_scheduler = true;
                 break;
             case BONG_TANGENT_SCHEDULER:
                 LOG_INFO("get_sigmas with bong_tangent scheduler");
@@ -618,6 +627,7 @@ struct Denoiser {
             case LCM_SCHEDULER:
                 LOG_INFO("get_sigmas with LCM scheduler");
                 scheduler = std::make_shared<LCMScheduler>();
+                shifted_scheduler = true;
                 break;
             case LTX2_SCHEDULER:
                 LOG_INFO("get_sigmas with LTX2 scheduler");
@@ -628,7 +638,14 @@ struct Denoiser {
                 scheduler = std::make_shared<DiscreteScheduler>();
                 break;
         }
-        return scheduler->get_sigmas(n, sigma_min(), sigma_max(), bound_t_to_sigma);
+        std::vector<float> sigmas = scheduler->get_sigmas(n, sigma_min(), sigma_max(), bound_t_to_sigma);
+
+        if (!shifted_scheduler) {
+            for (size_t i = 0; i < sigmas.size() - 1; ++i) {
+                sigmas[i] = apply_shift_warp(sigmas[i]);
+            }
+        }
+        return sigmas;
     }
 };
 
@@ -763,6 +780,10 @@ struct DiscreteFlowDenoiser : public Denoiser {
     float sigma_to_t(float sigma) override {
         return sigma * 1000.f;
     }
+    
+    float apply_shift_warp(float sigma) override {
+        return time_snr_shift(shift, sigma);
+    }
 
     float t_to_sigma(float t) override {
         t = t + 1;
@@ -798,9 +819,13 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
         return sigma;
     }
 
+    float apply_shift_warp(float sigma) override {
+        return flux_time_shift(shift, 1.0f, sigma);
+    }
+
     float t_to_sigma(float t) override {
         t = t + 1;
-        return flux_time_shift(shift, 1.0f, t / TIMESTEPS);
+        return apply_shift_warp(t / TIMESTEPS);
     }
 };