54 lines
2.3 KiB
Python
54 lines
2.3 KiB
Python
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)
|