new sampler features

modules/processing_diffusers.py

@@ -1,4 +1,5 @@
 from types import SimpleNamespace
+import re
 import os
 import time
 import math
@@ -80,7 +81,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             return kwargs
         elif shared.opts.nan_skip:
             assert not torch.isnan(latents[..., 0, 0]).all(), f'NaN detected at step {step}: Skipping...'
-        if len(getattr(p, "ip_adapter_names", [])) > 0:
+        if len(getattr(p, 'ip_adapter_names', [])) > 0:
             ip_adapter_scales = list(p.ip_adapter_scales)
             ip_adapter_starts = list(p.ip_adapter_starts)
             ip_adapter_ends = list(p.ip_adapter_ends)
@@ -112,7 +113,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
     def task_specific_kwargs(model):
         task_args = {}
         is_img2img_model = bool('Zero123' in shared.sd_model.__class__.__name__)
-        if len(getattr(p, 'init_images' ,[])) > 0:
+        if len(getattr(p, 'init_images', [])) > 0:
             p.init_images = [p.convert('RGB') for p in p.init_images]
         if sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE and not is_img2img_model:
             p.ops.append('txt2img')
@@ -121,13 +122,13 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                     'width': 8 * math.ceil(p.width / 8),
                     'height': 8 * math.ceil(p.height / 8),
                 }
-        elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
+        elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
             p.ops.append('img2img')
             task_args = {
                 'image': p.init_images,
                 'strength': p.denoising_strength,
             }
-        elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images' ,[])) > 0:
+        elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0:
             p.ops.append('instruct')
             task_args = {
                 'width': 8 * math.ceil(p.width / 8) if hasattr(p, 'width') else None,
@@ -135,7 +136,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 'image': p.init_images,
                 'strength': p.denoising_strength,
             }
-        elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0:
+        elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
             p.ops.append('inpaint')
             width, height = resize_init_images(p)
             task_args = {
@@ -196,7 +197,23 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
         parser = 'Fixed attention'
         clip_skip = kwargs.pop("clip_skip", 1)
-        steps = kwargs.get("num_inference_steps", 1)
+
+        steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1']))
+        if 'timesteps' in possible:
+            try:
+                timesteps = re.split(',| ', shared.opts.schedulers_timesteps)
+                timesteps = [int(x) for x in timesteps if x.isdigit()]
+                # AYS SD15: [999, 850, 736, 645, 545, 455, 343, 233, 124, 24]
+                # AYS SDXL: [999, 845, 730, 587, 443, 310, 193, 116, 53, 13]
+
+                if len(timesteps) > 0:
+                    args['timesteps'] = timesteps
+                    p.steps = len(timesteps)
+                    p.timesteps = timesteps
+                    steps = p.steps
+                    shared.log.debug(f'Sampler: steps={len(timesteps)} timesteps={timesteps}')
+            except Exception as e:
+                shared.log.error(f'Sampler timesteps: {e}')
         if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__ and 'Onnx' not in model.__class__.__name__:
             try:
                 prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, steps=steps, clip_skip=clip_skip)
@@ -343,6 +360,11 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             sampler = sd_samplers.all_samplers_map.get(sampler_selection, None)
             if sampler is None:
                 sampler = sd_samplers.all_samplers_map.get("UniPC")
+            if len(getattr(p, 'timesteps', [])) > 0:
+                if 'schedulers_use_karras' in shared.opts.data:
+                    shared.opts.data['schedulers_use_karras'] = False
+                else:
+                    shared.opts.schedulers_use_karras = False
             sampler = sd_samplers.create_sampler(sampler.name, sd_model)
             sampler_options = []
             if sampler.config.get('use_karras_sigmas', False):
@@ -438,7 +460,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) # reset pipeline
     if hasattr(shared.sd_model, 'unet') and hasattr(shared.sd_model.unet, 'config') and hasattr(shared.sd_model.unet.config, 'in_channels') and shared.sd_model.unet.config.in_channels == 9 and not is_control:
         shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) # force pipeline
-        if len(getattr(p, 'init_images' ,[])) == 0:
+        if len(getattr(p, 'init_images', [])) == 0:
             p.init_images = [TF.to_pil_image(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))]
 
     sd_models.move_model(shared.sd_model, devices.device)
@@ -449,7 +471,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
 
     shared.sd_model = update_pipeline(shared.sd_model, p)
     shared.log.info(f'Base: class={shared.sd_model.__class__.__name__}')
-    update_sampler(shared.sd_model)
     base_args = set_pipeline_args(
         model=shared.sd_model,
         prompts=p.prompts,
@@ -466,6 +487,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         clip_skip=p.clip_skip,
         desc='Base',
     )
+    update_sampler(shared.sd_model)
     shared.state.sampling_steps = base_args.get('num_inference_steps', p.steps)
     p.extra_generation_params['Pipeline'] = shared.sd_model.__class__.__name__
     if shared.opts.scheduler_eta is not None and shared.opts.scheduler_eta > 0 and shared.opts.scheduler_eta < 1:
@@ -546,7 +568,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if p.hr_force:
             shared.state.job_count = 2 * p.n_iter
             shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
-            update_sampler(shared.sd_model, second_pass=True)
             shared.log.info(f'HiRes: class={shared.sd_model.__class__.__name__} sampler="{p.hr_sampler_name}"')
             if p.is_control and hasattr(p, 'task_args') and p.task_args.get('image', None) is not None:
                 if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
@@ -571,6 +592,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 strength=p.denoising_strength,
                 desc='Hires',
             )
+            update_sampler(shared.sd_model, second_pass=True)
             shared.state.job = 'hires'
             shared.state.sampling_steps = hires_args['num_inference_steps']
             try:
@@ -607,7 +629,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         sd_models_compile.openvino_recompile_model(p, hires=False, refiner=True)
         shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE)
         shared.sd_refiner = sd_models.set_diffuser_pipe(shared.sd_refiner, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)
-        update_sampler(shared.sd_refiner, second_pass=True)
         for i in range(len(output.images)):
             image = output.images[i]
             noise_level = round(350 * p.denoising_strength)
@@ -636,6 +657,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 clip_skip=p.clip_skip,
                 desc='Refiner',
             )
+            update_sampler(shared.sd_refiner, second_pass=True)
             shared.state.sampling_steps = refiner_args['num_inference_steps']
             try:
                 if 'requires_aesthetics_score' in shared.sd_refiner.config: # sdxl-model needs false and sdxl-refiner needs true

modules/processing_helpers.py

@@ -420,6 +420,8 @@ def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
     return prompts, negative_prompts, prompts_2, negative_prompts_2
 
 def calculate_base_steps(p, use_denoise_start, use_refiner_start):
+    if len(getattr(p, 'timesteps', [])) > 0:
+        return None
     is_txt2img = sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE
     if not is_txt2img:
         if use_denoise_start and shared.sd_model_type == 'sdxl':
@@ -436,6 +438,8 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
     return max(1, int(steps))
 
 def calculate_hires_steps(p):
+    if len(getattr(p, 'timesteps', [])) > 0:
+        return None
     if p.hr_second_pass_steps > 0:
         steps = (p.hr_second_pass_steps // p.denoising_strength) + 1
     elif p.denoising_strength > 0:
@@ -446,6 +450,8 @@ def calculate_hires_steps(p):
     return max(1, int(steps))
 
 def calculate_refiner_steps(p):
+    if len(getattr(p, 'timesteps', [])) > 0:
+        return None
     if "StableDiffusionXL" in shared.sd_refiner.__class__.__name__:
         if p.refiner_start > 0 and p.refiner_start < 1:
             #steps = p.refiner_steps // (1 - p.refiner_start) # SDXL with denoise strenght

modules/sd_samplers_diffusers.py

@@ -8,6 +8,7 @@ from modules.tcd import TCDScheduler
 debug = shared.log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug('Trace: SAMPLER')
 
+
 try:
     from diffusers import (
         CMStochasticIterativeScheduler,
@@ -35,6 +36,7 @@ except Exception as e:
     import diffusers
     shared.log.error(f'Diffusers import error: version={diffusers.__version__} error: {e}')
 
+
 config = {
     # beta_start, beta_end are typically per-scheduler, but we don't want them as they should be taken from the model itself as those are values model was trained on
     # prediction_type is ideally set in model as well, but it maybe needed that we do auto-detect of model type in the future

modules/shared.py

@@ -701,11 +701,12 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"),
     # managed from ui.py for backend diffusers
     "schedulers_sep_diffusers": OptionInfo("<h2>Diffusers specific config</h2>", "", gr.HTML),
     "schedulers_dpm_solver": OptionInfo("sde-dpmsolver++", "DPM solver algorithm", gr.Radio, {"choices": ['dpmsolver++', 'sde-dpmsolver++']}),
-    "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Radio, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2']}),
+    "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2']}),
     'schedulers_beta_start': OptionInfo(0, "Beta start", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}),
     'schedulers_beta_end': OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}),
-    "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Radio, {"choices": ['default', 'linspace', 'leading', 'trailing']}),
+    "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Dropdown, {"choices": ['default', 'linspace', 'leading', 'trailing']}),
     'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1}),
+    'schedulers_timesteps': OptionInfo('', "Timesteps"),
     "schedulers_rescale_betas": OptionInfo(False, "Rescale betas with zero terminal SNR", gr.Checkbox),
 
     # managed from ui.py for backend original k-diffusion

modules/ui_sections.py

@@ -221,8 +221,28 @@ def create_sampler_options(tabname):
         shared.opts.data['schedulers_use_thresholding'] = 'dynamic thresholding' in sampler_options
         shared.opts.data['schedulers_use_loworder'] = 'low order' in sampler_options
         shared.opts.data['schedulers_rescale_betas'] = 'rescale beta' in sampler_options
+        shared.log.debug(f'Sampler set options: {sampler_options}')
         shared.opts.save(shared.config_filename, silent=True)
 
+    def set_sampler_timesteps(spacing, timesteps):
+        shared.log.debug(f'Sampler set options: spacing={spacing} timesteps={timesteps}')
+        if 'schedulers_timestep_spacing' in shared.opts.data:
+            shared.opts.data['schedulers_timestep_spacing'] = spacing
+        else:
+            shared.opts.schedulers_timestep_spacing = spacing
+        if 'schedulers_timesteps' in shared.opts.data:
+            shared.opts.data['schedulers_timesteps'] = timesteps
+        else:
+            shared.opts.schedulers_timesteps = timesteps
+        shared.opts.save(shared.config_filename, silent=True)
+
+    def set_sampler_preset(preset):
+        if preset == 'AYS SD15':
+            return '999,850,736,645,545,455,343,233,124,24'
+        if preset == 'AYS SDXL':
+            return '999,845,730,587,443,310,193,116,53,13'
+        return ''
+
     if shared.backend == shared.Backend.ORIGINAL:
         with gr.Row(elem_classes=['flex-break']):
             options = ['brownian noise', 'discard penultimate sigma']
@@ -244,7 +264,15 @@ def create_sampler_options(tabname):
             values += ['low order'] if shared.opts.data.get('schedulers_use_loworder', True) else []
             values += ['rescale beta'] if shared.opts.data.get('schedulers_rescale_betas', False) else []
             sampler_options = gr.CheckboxGroup(label='Sampler options', elem_id=f"{tabname}_sampler_options", choices=options, value=values, type='value')
+        with gr.Row(elem_classes=['flex-break']):
+            sampler_spacing = gr.Dropdown(label='Timestep spacing', elem_id=f"{tabname}_sampler_spacing", choices=['default', 'linspace', 'leading', 'trailing'], value=shared.opts.schedulers_timestep_spacing, type='value')
+            sampler_presets = gr.Dropdown(label='Presets', elem_id=f"{tabname}_sampler_presets", choices=['None', 'AYS SD15', 'AYS SDXL'], value='None', type='value')
+        with gr.Row(elem_classes=['flex-break']):
+            sampler_timesteps = gr.Textbox(label='Timesteps override', elem_id=f"{tabname}_sampler_timesteps", value=shared.opts.schedulers_timesteps)
         sampler_options.change(fn=set_sampler_diffuser_options, inputs=[sampler_options], outputs=[])
+        sampler_spacing.change(fn=set_sampler_timesteps, inputs=[sampler_spacing, sampler_timesteps], outputs=[])
+        sampler_timesteps.change(fn=set_sampler_timesteps, inputs=[sampler_spacing, sampler_timesteps], outputs=[])
+        sampler_presets.change(fn=set_sampler_preset, inputs=[sampler_presets], outputs=[sampler_timesteps])
 
 
 def create_hires_inputs(tab):