samplers add manual sigma adjustment

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -107,7 +107,8 @@ Pretty big performance updates to a) Any model using DiT based architecture: new
   - update `diffusers` and other requirements  
   - rename vae, unet and text-encoder settings *None* to *Default* to avoid confusion  
   - **CLI**: add `cli/api-grid.py` which can generate grids using params-from-file for x/y axis  
-  - LoRA enable memory cache by default  
+  - **LoRA** enable memory cache by default  
+  - **Samplers** add ability to set sigma adjustment for each sampler  
 - **Wiki/Docs**  
   - updated [Models](https://github.com/vladmandic/sdnext/wiki/Models) info  
   - new [Video](https://github.com/vladmandic/sdnext/wiki/Video) guide  

modules/processing_callbacks.py

@@ -59,8 +59,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
     if debug:
         debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}')
     shared.state.step()
-    # order = getattr(pipe.scheduler, "order", 1) if hasattr(pipe, 'scheduler') else 1
-    # shared.state.sampling_step = step // order
     if shared.state.interrupted or shared.state.skipped:
         raise AssertionError('Interrupted...')
     if shared.state.paused:
@@ -125,6 +123,8 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
             try:
                 shared.state.current_sigma = pipe.scheduler.sigmas[pipe.scheduler.step_index-1]
                 shared.state.current_sigma_next = pipe.scheduler.sigmas[pipe.scheduler.step_index]
+                if (shared.opts.schedulers_sigma_adjust != 1.0) and (timestep > 1000 * shared.opts.schedulers_sigma_adjust_min) and (timestep < 1000 * shared.opts.schedulers_sigma_adjust_max):
+                    pipe.scheduler.sigmas[pipe.scheduler.step_index+1] = pipe.scheduler.sigmas[pipe.scheduler.step_index+1] * shared.opts.schedulers_sigma_adjust
             except Exception:
                 pass
     except Exception as e:

modules/schedulers/scheduler_dpm_flowmatch.py

@@ -509,7 +509,10 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         def t_fn(_sigma: torch.Tensor) -> torch.Tensor:
             return _sigma.log().neg()
         sigma = self.sigmas[self.step_index]
-        sigma_next = self.sigmas[self.step_index + 1]
+        try:
+            sigma_next = self.sigmas[self.step_index + 1]
+        except Exception:
+            sigma_next = self.sigmas[-1]
         sigma_prev = self.sigmas[self.step_index - 1]
         if self.config.algorithm_type == "dpmsolver2":
             if self.config.solver_order == 2:

modules/sd_samplers_diffusers.py

@@ -193,8 +193,7 @@ class DiffusionSampler:
         self.name = name
         self.config = {}
         self.sampler = None
-        # if not hasattr(model, 'scheduler'):
-        #    return
+
         if getattr(model, "default_scheduler", None) is None and (model is not None): # sanity check
             model.default_scheduler = copy.deepcopy(model.scheduler)
         for key, value in config.get('All', {}).items(): # apply global defaults
@@ -217,6 +216,7 @@ class DiffusionSampler:
         for key, value in kwargs.items(): # apply user args, if any
             if key in self.config:
                 self.config[key] = value
+
         # finally apply user preferences
         if shared.opts.schedulers_prediction_type != 'default':
             self.config['prediction_type'] = shared.opts.schedulers_prediction_type
@@ -283,6 +283,7 @@ class DiffusionSampler:
             del self.config['prediction_type']
         if 'SGM' in name:
             self.config['timestep_spacing'] = 'trailing'
+
         # validate all config params
         signature = inspect.signature(constructor, follow_wrapped=True)
         possible = signature.parameters.keys()
@@ -293,7 +294,8 @@ class DiffusionSampler:
         debug_log(f'Sampler: name="{name}"')
         debug_log(f'Sampler: config={self.config}')
         debug_log(f'Sampler: signature={possible}')
-        # shared.log.debug_log(f'Sampler: sampler="{name}" config={self.config}')
+
+        # finally create the new sampler
         try:
             sampler = constructor(**self.config)
         except Exception as e:
@@ -302,21 +304,25 @@ class DiffusionSampler:
                 errors.display(e, 'Samplers')
             self.sampler = None
             return
-        accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
-        accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
-        accept_scale_noise = hasattr(sampler, "scale_noise")
-        debug_log(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
-        if ('StableDiffusion3' in model.__class__.__name__) and (not accept_scale_noise):
-            shared.log.warning(f'Sampler: sampler="{name}" does not implement scale noise')
-            self.sampler = None
-            return
+
+        if hasattr(sampler, 'set_timesteps'):
+            accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
+            accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys())
+            accept_scale_noise = hasattr(sampler, "scale_noise")
+            debug_log(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
+            if ('StableDiffusion3' in model.__class__.__name__) and (not accept_scale_noise):
+                shared.log.warning(f'Sampler: sampler="{name}" does not implement scale noise')
+                self.sampler = None
+                return
+
         self.sampler = sampler
         if name == 'DC Solver':
             if not hasattr(self.sampler, 'dc_ratios'):
                 pass
+
         # shared.log.debug_log(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}')
         self.sampler.name = name

modules/shared.py

@@ -792,6 +792,9 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"),
     'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1, "visible": native}),
     'schedulers_shift': OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}),
     'schedulers_dynamic_shift': OptionInfo(False, "Sampler dynamic shift", gr.Checkbox, {"visible": False}),
+    'schedulers_sigma_adjust': OptionInfo(1.0, "Sigma adjust", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01, "visible": False}),
+    'schedulers_sigma_adjust_min': OptionInfo(0.2, "Sigma adjust start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}),
+    'schedulers_sigma_adjust_max': OptionInfo(0.8, "Sigma adjust end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}),
 
     # managed from ui.py for backend original k-diffusion
     "always_batch_cond_uncond": OptionInfo(False, "Disable conditional batching", gr.Checkbox, {"visible": not native}),

modules/ui_sections.py

@@ -140,22 +140,22 @@ def create_seed_inputs(tab, reuse_visible=True, accordion=True, subseed_visible=
     return seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w
 
 
-def create_video_inputs(tab:str):
+def create_video_inputs(tab:str, show_always:bool=False):
     def video_type_change(video_type):
         return [
-            gr.update(visible=video_type != 'None'),
-            gr.update(visible=video_type in ['GIF', 'PNG']),
-            gr.update(visible=video_type not in ['None', 'GIF', 'PNG']),
-            gr.update(visible=video_type not in ['None', 'GIF', 'PNG']),
+            gr.update(visible=video_type != 'None' or show_always),
+            gr.update(visible=video_type in ['GIF', 'PNG'] or show_always),
+            gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always),
+            gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always),
         ]
     with gr.Column():
         video_codecs = ['None', 'GIF', 'PNG', 'MP4/MP4V', 'MP4/AVC1', 'MP4/JVT3', 'MKV/H264', 'AVI/DIVX', 'AVI/RGBA', 'MJPEG/MJPG', 'MPG/MPG1', 'AVR/AVR1']
         video_type = gr.Dropdown(label='Save video', choices=video_codecs, value='None', elem_id=f"{tab}_video_type")
     with gr.Column():
-        video_duration = gr.Slider(label='Duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=False, elem_id=f"{tab}_video_duration")
-        video_loop = gr.Checkbox(label='Loop', value=True, visible=False, elem_id=f"{tab}_video_loop")
-        video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=False, elem_id=f"{tab}_video_pad")
-        video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=False, elem_id=f"{tab}_video_interpolate")
+        video_duration = gr.Slider(label='Duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=show_always, elem_id=f"{tab}_video_duration")
+        video_loop = gr.Checkbox(label='Loop', value=True, visible=show_always, elem_id=f"{tab}_video_loop")
+        video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=show_always, elem_id=f"{tab}_video_pad")
+        video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=show_always, elem_id=f"{tab}_video_interpolate")
     video_type.change(fn=video_type_change, inputs=[video_type], outputs=[video_duration, video_loop, video_pad, video_interpolate])
     return video_type, video_duration, video_loop, video_pad, video_interpolate
 
@@ -274,6 +274,13 @@ def create_sampler_options(tabname):
         shared.opts.schedulers_shift = sampler_shift
         shared.opts.save(shared.config_filename, silent=True)
 
+    def set_sigma_ajust(val, start, end):
+        shared.log.debug(f'Sampler set options: sigma={val} min={start} max={end}')
+        shared.opts.schedulers_sigma_adjust = val
+        shared.opts.schedulers_sigma_adjust_min = start
+        shared.opts.schedulers_sigma_adjust_max = end
+        shared.opts.save(shared.config_filename, silent=True)
+
     # 'linear', 'scaled_linear', 'squaredcos_cap_v2'
     def set_sampler_preset(preset):
         if preset == 'AYS SD15':
@@ -305,6 +312,10 @@ def create_sampler_options(tabname):
         with gr.Row(elem_classes=['flex-break']):
             sampler_presets = gr.Dropdown(label='Timesteps presets', elem_id=f"{tabname}_sampler_presets", choices=['None', 'AYS SD15', 'AYS SDXL'], value='None', type='value')
             sampler_timesteps = gr.Textbox(label='Timesteps override', elem_id=f"{tabname}_sampler_timesteps", value=shared.opts.schedulers_timesteps)
+        with gr.Row(elem_classes=['flex-break']):
+            sampler_sigma_adjust_val = gr.Slider(minimum=0.5, maximum=1.5, step=0.01, label='Sigma adjust', value=shared.opts.schedulers_sigma_adjust, elem_id=f"{tabname}_sampler_sigma_adjust")
+            sampler_sigma_adjust_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Adjust start', value=shared.opts.schedulers_sigma_adjust_min, elem_id=f"{tabname}_sampler_sigma_adjust_min")
+            sampler_sigma_adjust_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Adjust end', value=shared.opts.schedulers_sigma_adjust_max, elem_id=f"{tabname}_sampler_sigma_adjust_max")
         with gr.Row(elem_classes=['flex-break']):
             sampler_order = gr.Slider(minimum=0, maximum=5, step=1, label="Sampler order", value=shared.opts.schedulers_solver_order, elem_id=f"{tabname}_sampler_order")
             sampler_shift = gr.Slider(minimum=0, maximum=10, step=0.1, label="Flow shift", value=shared.opts.schedulers_shift, elem_id=f"{tabname}_sampler_shift")
@@ -326,6 +337,9 @@ def create_sampler_options(tabname):
         sampler_order.change(fn=set_sampler_order, inputs=[sampler_order], outputs=[])
         sampler_shift.change(fn=set_sampler_shift, inputs=[sampler_shift], outputs=[])
         sampler_options.change(fn=set_sampler_options, inputs=[sampler_options], outputs=[])
+        sampler_sigma_adjust_val.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[])
+        sampler_sigma_adjust_min.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[])
+        sampler_sigma_adjust_max.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[])
 
 
 def create_hires_inputs(tab):

modules/ui_video.py

@@ -115,7 +115,7 @@ def create_ui():
                     with gr.Row():
                         save_frames = gr.Checkbox(label='Save image frames', value=False, elem_id="video_save_frames")
                     with gr.Row():
-                        video_type, video_duration, video_loop, video_pad, video_interpolate = ui_sections.create_video_inputs(tab='video')
+                        video_type, video_duration, video_loop, video_pad, video_interpolate = ui_sections.create_video_inputs(tab='video', show_always=True)
                 override_settings = ui_common.create_override_inputs('video')
 
             # output panel with gallery and video tabs

scripts/xyz_grid_classes.py

@@ -119,6 +119,7 @@ axis_options = [
     AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
     AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]),
     AxisOption("[Sampler] Sigma method", str, apply_setting("schedulers_sigma"), choices=lambda: ['default', 'karras', 'betas', 'exponential', 'lambdas']),
+    AxisOption("[Sampler] Sigma adjust", float, apply_setting("schedulers_sigma_adjust")),
     AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']),
     AxisOption("[Sampler] Timestep range", int, apply_setting("schedulers_timesteps_range")),
     AxisOption("[Sampler] Solver order", int, apply_setting("schedulers_solver_order")),