diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py index d5f0a49b..1443c5cd 100644 --- a/modules/generation_parameters_copypaste.py +++ b/modules/generation_parameters_copypaste.py @@ -306,6 +306,18 @@ Steps: 20, Sampler: Euler a, CFG scale: 7, Seed: 965400086, Size: 512x512, Model if "RNG" not in res: res["RNG"] = "GPU" + if "KDiff Schedule Type" not in res: + res["KDiff Schedule Type"] = "Automatic" + + if "KDiff Schedule max sigma" not in res: + res["KDiff Schedule max sigma"] = 14.6 + + if "KDiff Schedule min sigma" not in res: + res["KDiff Schedule min sigma"] = 0.3 + + if "KDiff Schedule rho" not in res: + res["KDiff Schedule rho"] = 7.0 + return res @@ -318,6 +330,10 @@ infotext_to_setting_name_mapping = [ ('Conditional mask weight', 'inpainting_mask_weight'), ('Model hash', 'sd_model_checkpoint'), ('ENSD', 'eta_noise_seed_delta'), + ('KDiff Schedule Type', 'k_sched_type'), + ('KDiff Schedule max sigma', 'sigma_max'), + ('KDiff Schedule min sigma', 'sigma_min'), + ('KDiff Schedule rho', 'rho'), ('Noise multiplier', 'initial_noise_multiplier'), ('Eta', 'eta_ancestral'), ('Eta DDIM', 'eta_ddim'), diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 638e0ac9..9c9d9f17 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -44,6 +44,14 @@ sampler_extra_params = { 'sample_dpm_2': ['s_churn', 's_tmin', 's_tmax', 's_noise'], } +k_diffusion_samplers_map = {x.name: x for x in samplers_data_k_diffusion} +k_diffusion_scheduler = { + 'Automatic': None, + 'karras': k_diffusion.sampling.get_sigmas_karras, + 'exponential': k_diffusion.sampling.get_sigmas_exponential, + 'polyexponential': k_diffusion.sampling.get_sigmas_polyexponential +} + class CFGDenoiser(torch.nn.Module): """ @@ -265,6 +273,13 @@ class KDiffusionSampler: try: return func() + except RecursionError: + print( + 'Encountered RecursionError during sampling, returning last latent. ' + 'rho >5 with a polyexponential scheduler may cause this error. ' + 'You should try to use a smaller rho value instead.' + ) + return self.last_latent except sd_samplers_common.InterruptedException: return self.last_latent @@ -304,6 +319,29 @@ class KDiffusionSampler: if p.sampler_noise_scheduler_override: sigmas = p.sampler_noise_scheduler_override(steps) + elif opts.k_sched_type != "Automatic": + m_sigma_min, m_sigma_max = (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item()) + sigma_min, sigma_max = (0.1, 10) + sigmas_kwargs = { + 'sigma_min': sigma_min if opts.use_old_karras_scheduler_sigmas else m_sigma_min, + 'sigma_max': sigma_max if opts.use_old_karras_scheduler_sigmas else m_sigma_max + } + + sigmas_func = k_diffusion_scheduler[opts.k_sched_type] + p.extra_generation_params["KDiff Schedule Type"] = opts.k_sched_type + + if opts.sigma_min != 0.3: + # take 0.0 as model default + sigmas_kwargs['sigma_min'] = opts.sigma_min or m_sigma_min + p.extra_generation_params["KDiff Schedule min sigma"] = opts.sigma_min + if opts.sigma_max != 14.6: + sigmas_kwargs['sigma_max'] = opts.sigma_max or m_sigma_max + p.extra_generation_params["KDiff Schedule max sigma"] = opts.sigma_max + if opts.k_sched_type != 'exponential': + sigmas_kwargs['rho'] = opts.rho + p.extra_generation_params["KDiff Schedule rho"] = opts.rho + + sigmas = sigmas_func(n=steps, **sigmas_kwargs, device=shared.device) elif self.config is not None and self.config.options.get('scheduler', None) == 'karras': sigma_min, sigma_max = (0.1, 10) if opts.use_old_karras_scheduler_sigmas else (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item()) diff --git a/modules/shared.py b/modules/shared.py index a5e7824a..364a5991 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -518,6 +518,10 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters" 's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), 's_tmin': OptionInfo(0.0, "sigma tmin", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), 's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), + 'k_sched_type': OptionInfo("Automatic", "scheduler type", gr.Dropdown, {"choices": ["Automatic", "karras", "exponential", "polyexponential"]}), + 'sigma_max': OptionInfo(14.6, "sigma max", gr.Number).info("the maximum noise strength for the scheduler. Set to 0 to use the same value which 'xxx karras' samplers use."), + 'sigma_min': OptionInfo(0.3, "sigma min", gr.Number).info("the minimum noise strength for the scheduler. Set to 0 to use the same value which 'xxx karras' samplers use."), + 'rho': OptionInfo(7.0, "rho", gr.Number).info("higher will make a more steep noise scheduler (decrease faster). default for karras is 7.0, for polyexponential is 1.0"), 'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}).info("ENSD; does not improve anything, just produces different results for ancestral samplers - only useful for reproducing images"), 'always_discard_next_to_last_sigma': OptionInfo(False, "Always discard next-to-last sigma").link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/6044"), 'uni_pc_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"]}), diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index da820b39..089d375e 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -10,7 +10,7 @@ import numpy as np import modules.scripts as scripts import gradio as gr -from modules import images, sd_samplers, processing, sd_models, sd_vae +from modules import images, sd_samplers, processing, sd_models, sd_vae, sd_samplers_kdiffusion from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img from modules.shared import opts, state import modules.shared as shared @@ -220,6 +220,10 @@ axis_options = [ AxisOption("Sigma min", float, apply_field("s_tmin")), AxisOption("Sigma max", float, apply_field("s_tmax")), AxisOption("Sigma noise", float, apply_field("s_noise")), + AxisOption("KDiff Schedule Type", str, apply_override("k_sched_type"), choices=lambda: list(sd_samplers_kdiffusion.k_diffusion_scheduler)), + AxisOption("KDiff Schedule min sigma", float, apply_override("sigma_min")), + AxisOption("KDiff Schedule max sigma", float, apply_override("sigma_max")), + AxisOption("KDiff Schedule rho", float, apply_override("rho")), AxisOption("Eta", float, apply_field("eta")), AxisOption("Clip skip", int, apply_clip_skip), AxisOption("Denoising", float, apply_field("denoising_strength")),