Alternate sampler noise schedules for stable-diffusion-webui (original) (raw)

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)