From 964b4c9e5a6f7bdab943007edeeb3b2c828ac158 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Thu, 28 Nov 2024 09:11:42 -0500
Subject: [PATCH] euler flowmatch add sigma methods

Signed-off-by: Vladimir Mandic <mandic00@live.com>
---
 CHANGELOG.md                     |  3 ++-
 modules/sd_samplers_diffusers.py | 20 ++++++++++----------
 2 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b3f2282c0..f88c78302 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,7 +50,8 @@
   - control: add stats
   - browser->server logging framework  
 - **Sampler** improvements  
-  - update DPM FlowMatch samplers  
+  - Euler FlowMatch: add sigma methods (*karras/exponential/betas*)  
+  - DPM FlowMatch: update all and add sigma methods  
 
 ### Fixes  
 
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index 9f24d5a91..4672df92e 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/modules/sd_samplers_diffusers.py
@@ -69,7 +69,7 @@ config = {
     'Euler a': { 'steps_offset': 0, 'rescale_betas_zero_snr': False, 'timestep_spacing': 'linspace' },
     'Euler SGM': { 'steps_offset': 0, 'interpolation_type': "linear", 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'trailing', 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_karras_sigmas': False, 'prediction_type': "sample" },
     'Euler EDM': { 'sigma_schedule': "karras" },
-    'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, 'use_dynamic_shifting': False },
+    'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, 'use_dynamic_shifting': False, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False },
 
     'DPM++': { 'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'final_sigmas_type': 'sigma_min' },
     'DPM++ 1S': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'use_lu_lambdas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 1 },
@@ -200,16 +200,16 @@ class DiffusionSampler:
         timesteps = re.split(',| ', shared.opts.schedulers_timesteps)
         timesteps = [int(x) for x in timesteps if x.isdigit()]
         if len(timesteps) == 0:
-            if 'use_beta_sigmas' in self.config:
-                self.config['use_beta_sigmas'] = shared.opts.schedulers_sigma == 'beta'
-            if 'use_karras_sigmas' in self.config:
-                self.config['use_karras_sigmas'] = shared.opts.schedulers_sigma == 'karras'
-            if 'use_exponential_sigmas' in self.config:
-                self.config['use_exponential_sigmas'] = shared.opts.schedulers_sigma == 'exponential'
-            if 'use_lu_lambdas' in self.config:
-                self.config['use_lu_lambdas'] = shared.opts.schedulers_sigma == 'lambdas'
             if 'sigma_schedule' in self.config:
                 self.config['sigma_schedule'] = shared.opts.schedulers_sigma if shared.opts.schedulers_sigma != 'default' else None
+            if shared.opts.schedulers_sigma == 'betas' and 'use_beta_sigmas' in self.config:
+                self.config['use_beta_sigmas'] = True
+            elif shared.opts.schedulers_sigma == 'karras' and 'use_karras_sigmas' in self.config:
+                self.config['use_karras_sigmas'] = True
+            elif shared.opts.schedulers_sigma == 'exponential' and 'use_exponential_sigmas' in self.config:
+                self.config['use_exponential_sigmas'] = True
+            elif shared.opts.schedulers_sigma == 'lambdas' and 'use_lu_lambdas' in self.config:
+                self.config['use_lu_lambdas'] = True
         else:
             pass # timesteps are set using set_timesteps in set_pipeline_args
 
@@ -236,7 +236,7 @@ class DiffusionSampler:
         if 'use_dynamic_shifting' in self.config:
             if 'Flux' in model.__class__.__name__:
                 self.config['use_dynamic_shifting'] = shared.opts.schedulers_dynamic_shift
-        if 'use_beta_sigmas' in self.config:
+        if 'use_beta_sigmas' in self.config and 'sigma_schedule' in self.config:
             self.config['use_beta_sigmas'] = 'StableDiffusion3' in model.__class__.__name__
         if 'rescale_betas_zero_snr' in self.config:
             self.config['rescale_betas_zero_snr'] = shared.opts.schedulers_rescale_betas