From 8397e5ef2df63250988d06c4a64bf608a46b3634 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 30 Dec 2024 12:21:47 -0500
Subject: [PATCH] samplers autodetect if sigma is required and available

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5e11c04c5..06524e1c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,12 +9,16 @@
   - add legacy option in *settings -> networks*  
 - **HunyuanVideo** optimizations: full offload, quantization and tiling support  
 - **LTXVideo** optimizations: full offload, quantization and tiling support  
-- VAE tiling granular options in *settings -> variable auto encoder*  
-- UI: live preview optimizations and error handling  
-- UI: live preview sigma calulations, thanks @Disty0  
-- UI: CSS optimizations when log view is disabled  
-- Samplers: add flow shift options and separate dynamic thresholding from dynamic shifting  
-- Refactor: remove all LDM imports if running in native mode  
+- **VAE**: tiling granular options in *settings -> variable auto encoder*  
+- **UI**:
+  - live preview optimizations and error handling  
+  - live preview high quality for flow models, thanks @Disty0  
+  - CSS optimizations when log view is disabled  
+- **Samplers**:
+  - add flow shift options and separate dynamic thresholding from dynamic shifting  
+  - autodetect matching sigma capabilities  
+- **Refactor**:
+  - remove all LDM imports if running in native mode  
 - **Fixes**  
   - do not show disabled networks  
   - image width/height calculation when doing img2img  
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 252e52d0f..89c4e5edf 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -92,9 +92,12 @@ def create_sampler(name, model):
         sampler = config.constructor(model)
         if sampler is None:
             sampler = config.constructor(model)
+        if sampler is None or sampler.sampler is None:
+            model.scheduler = copy.deepcopy(model.default_scheduler)
+        else:
+            model.scheduler = sampler.sampler
         if not hasattr(model, 'scheduler_config'):
-            model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler.sampler, 'config') else {}
-        model.scheduler = sampler.sampler
+            model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler, 'sampler') and hasattr(sampler.sampler, 'config') else {}
         if hasattr(model, "prior_pipe") and hasattr(model.prior_pipe, "scheduler"):
             model.prior_pipe.scheduler = sampler.sampler
             model.prior_pipe.scheduler.config.clip_sample = False
@@ -102,8 +105,9 @@ def create_sampler(name, model):
             shared.state.prediction_type = "flow_prediction"
         elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"):
             shared.state.prediction_type = model.scheduler.config.prediction_type
-        clean_config = {k: v for k, v in sampler.config.items() if v is not None and v is not False}
-        shared.log.debug(f'Sampler: "{sampler.name}" class={model.scheduler.__class__.__name__} config={clean_config}')
+        clean_config = {k: v for k, v in model.scheduler.config.items() if not k.startswith('_') and v is not None and v is not False}
+        name = sampler.name if sampler is not None and sampler.sampler is not None else 'Default'
+        shared.log.debug(f'Sampler: "{name}" class={model.scheduler.__class__.__name__} config={clean_config}')
         return sampler.sampler
     else:
         return None
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index 13a5c3ad9..06459d8a8 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/modules/sd_samplers_diffusers.py
@@ -268,7 +268,15 @@ class DiffusionSampler:
         debug(f'Sampler: config={self.config}')
         debug(f'Sampler: signature={possible}')
         # shared.log.debug(f'Sampler: sampler="{name}" config={self.config}')
-        self.sampler = constructor(**self.config)
+        sampler = constructor(**self.config)
+        accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
+        accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
+        debug(f'Sampler: sampler="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}')
+        if ('Flux' in model.__class__.__name__) and (not accept_sigmas):
+            shared.log.warning(f'Sampler: sampler="{name}" does not accept sigmas')
+            self.sampler = None
+            return
+        self.sampler = sampler
         if name == 'DC Solver':
             if not hasattr(self.sampler, 'dc_ratios'):
                 pass