new sampler features

CHANGELOG.md

@@ -1,25 +1,5 @@
 # Change Log for SD.Next
 
-## Pending
-
-### requires `diffusers-0.28.0.dev0`
-
-- PixArt-Σ
-- IP adapter masking
-- InstantStyle
-
-### Note
-
-This version of [SD.Next](https://github.com/vladmandic/automatic) ships with a preview of the new [ModernUI](https://github.com/BinaryQuantumSoul/sdnext-modernui)  
-For details on how to enable and use it, see [Home](https://github.com/BinaryQuantumSoul/sdnext-modernui) and [WiKi](https://github.com/vladmandic/automatic/wiki/Themes)  
-**ModernUI** is still in early development and not all features are available yet, please report [issues and feedback](https://github.com/BinaryQuantumSoul/sdnext-modernui/issues)  
-Thanks to @BinaryQuantumSoul for his hard work on this project!
-
-*Note*: [SD.Next](https://github.com/vladmandic/automatic) is no longer marked as a fork of [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui/) and github project has been fully detached  
-Given huge number of changes with *+3443/-3342* commits diff (at the time of fork detach) over the past year,  
-a completely different backend/engine and a change of focus,  
-it is time to give credit to original [author](https://github.com/auTOMATIC1111),  and move on!  
-
 ## Update for 2024-05-11
 
 - **Features**:
@@ -135,7 +115,12 @@ it is time to give credit to original [author](https://github.com/auTOMATIC1111)
   - Add *Euler SGM* variation (e.g. SGM Uniform), optimized for SDXL-Lightning models  
     *note*: you can use other samplers as well with SDXL-Lightning models  
   - Add *CMSI* sampler, optimized for consistency models  
-  - Add option *timestep spacing* to sampler settings  
+  - Add option *timestep spacing* to sampler settings and sampler section in main ui
+    Note: changing timestep spacing changes behavior of sampler and can help to make any sampler turbo/lightning compatibile
+  - Add option *timesteps* to manually set timesteps instead of relying on steps+spacing  
+    Additionally, presets from nVidia's align-you-steps reasearch are provided  
+    Result is that perfectly aligned steps can drastically reduce number of steps needed!  
+    For example, AIY preset alows DPM++2M to run in ~10 steps with quality equallying ~30 steps!  
 - **IPEX**, thanks @Disty0
   - Update to *IPEX 2.1.20* on Linux  
     requires removing the venv folder to update properly  

TODO.md

@@ -4,10 +4,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 ## Candidates for next release
 
-### Logistics
-
-- integrate uiux
-
 ### Models & Pipelines
 
 - stable diffusion 3.0
@@ -26,8 +22,45 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 - control api scripts compatibility
 
-### Defer
+## Future release notes
 
-- ella: <https://github.com/TencentQQGYLab/ELLA>
-- x-adapter: <https://github.com/showlab/X-Adapter>
-- samplers: [smea](https://github.com/Koishi-Star/Euler-Smea-Dyn-Sampler), [restart](https://github.com/Newbeeer/diffusion_restart_sampling), [vdm](https://github.com/huggingface/diffusers/pull/7737)
+### requires `diffusers-0.28.0.dev0`
+
+- PixArt-Σ
+- IP adapter masking
+- InstantStyle
+- Sampler timesteps
+
+## Note
+
+New release has been baking in `dev` for a longer than usual, but changes are massive...
+
+First, this version of [SD.Next](https://github.com/vladmandic/automatic) ships with a preview of the new [ModernUI](https://github.com/BinaryQuantumSoul/sdnext-modernui)  
+For details on how to enable and use it, see [Home](https://github.com/BinaryQuantumSoul/sdnext-modernui) and [WiKi](https://github.com/vladmandic/automatic/wiki/Themes)  
+**ModernUI** is still in early development and not all features are available yet, please report [issues and feedback](https://github.com/BinaryQuantumSoul/sdnext-modernui/issues)  
+Thanks to @BinaryQuantumSoul for his hard work on this project!  
+
+*What else?*
+
+### New built-in features
+
+- HiDiffusion allows generating very-high resolution images out-of-the-box using standard models  
+- Perturbed-Attention Guidance (PAG) enhances sample quality in addition to standard CFG scale  
+- IP adapter masking allows to use multiple input images for each segment of the input image  
+- IP adapter InstantStyle implementation  
+- Token Downsampling (ToDO) provides significant speedups with minimal-to-none quality loss  
+- Samplers optimizations that allow normal samplers to complete work in 1/3 of the steps!  
+  Yup, even popular DPM++2M can now run in 10 steps with quality equaling 30 steps  
+- Better outpainting
+
+### New models
+
+While still waiting for *Stable Diffusion 3.0*, there have been some significant models released in the meantime:
+- [PixArt-Σ](https://pixart-alpha.github.io/PixArt-sigma-project/), high end diffusion Transformer model (DiT) with a T5 encoder/decoder capable of directly generating images at 4K resolution  
+- [SDXS](https://github.com/IDKiro/sdxs), an extremely fast 1-step generation consistency model that also uses TAESD as quick VAE out-of-the-box  
+- [Hyper-SD](https://huggingface.co/ByteDance/Hyper-SD), 1-step, 2-step, 4-step and 8-step optimized models
+
+*Note*  
+[SD.Next](https://github.com/vladmandic/automatic) is no longer marked as a fork of [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui/) and github project has been fully detached  
+Given huge number of changes with *+3443/-3342* commits diff (at the time of fork detach) over the past year,  
+a completely different backend/engine and a change of focus, it is time to give credit to original [author](https://github.com/auTOMATIC1111),  and move on!  

javascript/sdnext.css

@@ -256,6 +256,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt
 #ui_defaults_review { margin: 1em; }
 .ar-dropdown { font-size: 0.9em; min-width: 5.5em !important; max-width: 5.5em !important; margin: 0 !important; padding: 0 !important; align-content: center; }
 .ar-dropdown div { margin: 0; background: var(--background-color)}
+#txt2img_sampler_timesteps, #img2img_sampler_timesteps { max-width: calc(var(--left-column) - 50px); }
 
 /* extras */
 .extras { gap: 0.2em 1em !important }

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):

scripts/xyz_grid.py

@@ -231,8 +231,6 @@ axis_options = [
     AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)),
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
     AxisOption("Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]),
-    AxisOptionTxt2Img("Sampler", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
-    AxisOptionImg2Img("Sampler", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]),
     AxisOption("Seed", int, apply_field("seed")),
     AxisOption("Steps", int, apply_field("steps")),
     AxisOption("CFG Scale", float, apply_field("cfg_scale")),
@@ -244,8 +242,9 @@ axis_options = [
     AxisOption("Prompt order", str_permutations, apply_order, fmt=format_value_join_list),
     AxisOption("Model dictionary", str, apply_dict, fmt=format_value, cost=1.0, choices=lambda: ['None'] + list(sd_models.checkpoints_list)),
     AxisOptionImg2Img("Image mask weight", float, apply_field("inpainting_mask_weight")),
-    AxisOption("[Postprocess] Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]),
-    AxisOption("[Postprocess] Face restore", str, apply_face_restore, fmt=format_value),
+    AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
+    AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]),
+    AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']),
     AxisOption("[Sampler] Sigma min", float, apply_field("s_min")),
     AxisOption("[Sampler] Sigma max", float, apply_field("s_max")),
     AxisOption("[Sampler] Sigma tmin", float, apply_field("s_tmin")),
@@ -263,6 +262,8 @@ axis_options = [
     AxisOption("[Refiner] Model", str, apply_refiner, fmt=format_value, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list)),
     AxisOption("[Refiner] Refiner start", float, apply_field("refiner_start")),
     AxisOption("[Refiner] Refiner steps", float, apply_field("refiner_steps")),
+    AxisOption("[Postprocess] Upscaler", str, apply_upscaler, choices=lambda: [x.name for x in shared.sd_upscalers][1:]),
+    AxisOption("[Postprocess] Face restore", str, apply_face_restore, fmt=format_value),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),