using dynamic shift Scheduler in flux2

diffsynth/diffusion/flow_match.py

@@ -122,3 +122,20 @@ class FlowMatchScheduler():
         b = base_shift - m * base_seq_len
         mu = image_seq_len * m + b
         return mu
+
+    def compute_empirical_mu(self, image_seq_len: int, num_steps: int) -> float:
+        a1, b1 = 8.73809524e-05, 1.89833333
+        a2, b2 = 0.00016927, 0.45666666
+
+        if image_seq_len > 4300:
+            mu = a2 * image_seq_len + b2
+            return float(mu)
+
+        m_200 = a2 * image_seq_len + b2
+        m_10 = a1 * image_seq_len + b1
+
+        a = (m_200 - m_10) / 190.0
+        b = m_200 - 200.0 * a
+        mu = a * num_steps + b
+
+        return float(mu)

diffsynth/pipelines/flux2_image.py

@@ -86,9 +86,6 @@ class Flux2ImagePipeline(BasePipeline):
         # Progress bar
         progress_bar_cmd = tqdm,
     ):
-        # Scheduler
-        self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength)
-        
         # Parameters
         inputs_posi = {
             "prompt": prompt,
@@ -106,6 +103,12 @@ class Flux2ImagePipeline(BasePipeline):
         for unit in self.units:
             inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega)
 
+        # using dynamic shift Scheduler
+        self.scheduler.exponential_shift = True
+        self.scheduler.sigma_min = 1 / num_inference_steps
+        mu = self.scheduler.compute_empirical_mu(inputs_shared["latents"].shape[1], num_inference_steps)
+        self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, exponential_shift_mu=mu)
+        
         # Denoise
         self.load_models_to_device(self.in_iteration_models)
         models = {name: getattr(self, name) for name in self.in_iteration_models}