Create alternate_sampler_noise_schedules.py

This commit is contained in:
C43H66N12O12S2 2022-10-03 12:49:17 +03:00 committed by AUTOMATIC1111
parent c30c06db20
commit c1a068ed0a

View file

@ -0,0 +1,53 @@
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)