| import inspect | |
| from modules.processing import Processed, process_images | |
| import gradio as gr | |
| import modules.scripts as scripts | |
| import k_diffusion.sampling | |
| import torch | |
| class Script(scripts.Script): | |
| def title(self): | |
| return "Alternate Sampler Noise Schedules" | |
| def ui(self, is_img2img): | |
| noise_scheduler = gr.Dropdown(label="Noise Scheduler", choices=['Default','Karras','Exponential', 'Variance Preserving'], value='Default', type="index") | |
| sched_smin = gr.Slider(value=0.1, label="Sigma min", minimum=0.0, maximum=100.0, step=0.5) | |
| sched_smax = gr.Slider(value=10.0, label="Sigma max", minimum=0.0, maximum=100.0, step=0.5) | |
| sched_rho = gr.Slider(value=7.0, label="Sigma rho (Karras only)", minimum=7.0, maximum=100.0, step=0.5) | |
| sched_beta_d = gr.Slider(value=19.9, label="Beta distribution (VP only)",minimum=0.0, maximum=40.0, step=0.5) | |
| sched_beta_min = gr.Slider(value=0.1, label="Beta min (VP only)", minimum=0.0, maximum=40.0, step=0.1) | |
| sched_eps_s = gr.Slider(value=0.001, label="Epsilon (VP only)", minimum=0.001, maximum=1.0, step=0.001) | |
| return [noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s] | |
| def run(self, p, noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s): | |
| noise_scheduler_func_name = ['-','get_sigmas_karras','get_sigmas_exponential','get_sigmas_vp'][noise_scheduler] | |
| base_params = { | |
| "sigma_min":sched_smin, | |
| "sigma_max":sched_smax, | |
| "rho":sched_rho, | |
| "beta_d":sched_beta_d, | |
| "beta_min":sched_beta_min, | |
| "eps_s":sched_eps_s, | |
| "device":"cuda" if torch.cuda.is_available() else "cpu" | |
| } | |
| if hasattr(k_diffusion.sampling,noise_scheduler_func_name): | |
| sigma_func = getattr(k_diffusion.sampling,noise_scheduler_func_name) | |
| sigma_func_kwargs = {} | |
| for k,v in base_params.items(): | |
| if k in inspect.signature(sigma_func).parameters: | |
| sigma_func_kwargs[k] = v | |
| def substitute_noise_scheduler(n): | |
| return sigma_func(n,**sigma_func_kwargs) | |
| p.sampler_noise_scheduler_override = substitute_noise_scheduler | |
| return process_images(p) |