|
|
from typing import Any |
|
|
|
|
|
import gradio as gr |
|
|
|
|
|
import modules.scripts as scripts |
|
|
from modules.processing import StableDiffusionProcessing |
|
|
from modules.script_callbacks import on_ui_settings, on_app_started |
|
|
from modules.shared import OptionInfo, opts |
|
|
|
|
|
import lib_es.const as consts |
|
|
from lib_es.xyz import xyz_support |
|
|
from lib_es.samplers import add_extra_samplers |
|
|
from lib_es.schedulers import add_schedulers |
|
|
|
|
|
from modules.script_callbacks import on_before_ui |
|
|
|
|
|
def early_init(): |
|
|
add_extra_samplers() |
|
|
add_schedulers() |
|
|
|
|
|
on_before_ui(early_init) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def from_setting_or_default(key: str, default: None | Any) -> None | Any: |
|
|
return opts.data.get(key, default) |
|
|
|
|
|
|
|
|
def on_change_update_setting(key: str, value: Any) -> None: |
|
|
opts.set(key, value) |
|
|
|
|
|
|
|
|
class ExtraSamplerExtension(scripts.Script): |
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
self.xyz_cache = {} |
|
|
xyz_support(self.xyz_cache) |
|
|
|
|
|
def title(self): |
|
|
return "Extra Samplers" |
|
|
|
|
|
def show(self, is_img2img): |
|
|
return scripts.AlwaysVisible |
|
|
|
|
|
def ui(self, is_img2img): |
|
|
with gr.Accordion(label="Extra Samplers", open=False): |
|
|
with gr.Accordion(label="Adaptive Progressive", open=False): |
|
|
gr.Markdown("Adaptive progressive sampler that automatically adjusts to different step counts. ") |
|
|
gr.Markdown( |
|
|
"Phase ends are automatically adjusted based on the total number of steps. These are approximations" |
|
|
) |
|
|
with gr.Row(): |
|
|
euler_a_end = gr.Slider( |
|
|
minimum=0.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=from_setting_or_default(consts.AP_EULER_A_END, 0.35), |
|
|
label="Euler A end", |
|
|
) |
|
|
dpm_2m_end = gr.Slider( |
|
|
minimum=0.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=from_setting_or_default(consts.AP_DPM_2M_END, 0.75), |
|
|
label="DPM++ 2M end", |
|
|
) |
|
|
with gr.Row(): |
|
|
ancestral_eta = gr.Slider( |
|
|
minimum=0.0, |
|
|
maximum=1.0, |
|
|
step=0.05, |
|
|
value=from_setting_or_default(consts.AP_ANCESTRAL_ETA, 0.4), |
|
|
label="Ancestral Eta", |
|
|
) |
|
|
detail_strength = gr.Slider( |
|
|
minimum=0.0, |
|
|
maximum=10.0, |
|
|
step=0.1, |
|
|
value=from_setting_or_default(consts.AP_DETAIL_STRENGTH, 1.5), |
|
|
label="Detail Strength", |
|
|
) |
|
|
|
|
|
euler_a_end.change( |
|
|
fn=lambda value: on_change_update_setting(consts.AP_EULER_A_END, value), inputs=[euler_a_end] |
|
|
) |
|
|
dpm_2m_end.change( |
|
|
fn=lambda value: on_change_update_setting(consts.AP_DPM_2M_END, value), inputs=[dpm_2m_end] |
|
|
) |
|
|
ancestral_eta.change( |
|
|
fn=lambda value: on_change_update_setting(consts.AP_ANCESTRAL_ETA, value), |
|
|
inputs=[ancestral_eta], |
|
|
) |
|
|
detail_strength.change( |
|
|
fn=lambda value: on_change_update_setting(consts.AP_DETAIL_STRENGTH, value), |
|
|
inputs=[detail_strength], |
|
|
) |
|
|
|
|
|
with gr.Accordion(label="Langevin Euler", open=False): |
|
|
langevin_strength = gr.Slider( |
|
|
minimum=0.0, |
|
|
maximum=0.5, |
|
|
step=0.01, |
|
|
value=from_setting_or_default(consts.LANGEVIN_STRENGTH, 0.1), |
|
|
label="Langevin Strength", |
|
|
info="Langevin strength for Langevin Euler sampler. Adjust to control the amount of noise.", |
|
|
) |
|
|
langevin_strength.change( |
|
|
fn=lambda value: on_change_update_setting(consts.LANGEVIN_STRENGTH, value), |
|
|
inputs=[langevin_strength], |
|
|
) |
|
|
|
|
|
with gr.Accordion(label="Gradient Estimation", open=False): |
|
|
use_adaptive_steps = from_setting_or_default(consts.GE_USE_ADAPTIVE_STEPS, False) |
|
|
|
|
|
adaptive_steps = gr.Checkbox( |
|
|
label="Use Adaptive Steps", |
|
|
value=use_adaptive_steps, |
|
|
info="Modify the number of steps based on the noise schedule.", |
|
|
) |
|
|
use_timestep_adaptive_gamma = gr.Checkbox( |
|
|
label="Timestep-Based Adaptive Gamma", |
|
|
value=from_setting_or_default(consts.GE_USE_TIMESTEP_ADAPTIVE_GAMMA, False), |
|
|
info="Adjust gamma during generation.", |
|
|
) |
|
|
gamma = gr.Slider( |
|
|
minimum=consts.GE_MIN_GAMMA, |
|
|
maximum=consts.GE_MAX_GAMMA, |
|
|
step=0.05, |
|
|
value=from_setting_or_default(consts.GE_GAMMA, consts.GE_DEFAULT_GAMMA), |
|
|
label="Gamma", |
|
|
info="Gamma value for gradient estimation. Higher values increase the amount of noise.", |
|
|
interactive=not use_adaptive_steps, |
|
|
) |
|
|
gamma_offset = gr.Slider( |
|
|
minimum=consts.GE_MIN_GAMMA_OFFSET, |
|
|
maximum=consts.GE_MAX_GAMMA_OFFSET, |
|
|
step=0.05, |
|
|
value=from_setting_or_default(consts.GE_GAMMA_OFFSET, consts.GE_DEFAULT_GAMMA_OFFSET), |
|
|
label="Gamma Offset", |
|
|
info="Offset to add to the calculated gamma when using adaptive steps.", |
|
|
interactive=use_adaptive_steps, |
|
|
) |
|
|
gamma.change(fn=lambda value: on_change_update_setting(consts.GE_GAMMA, value), inputs=[gamma]) |
|
|
gamma_offset.change( |
|
|
fn=lambda value: on_change_update_setting(consts.GE_GAMMA_OFFSET, value), inputs=[gamma_offset] |
|
|
) |
|
|
|
|
|
|
|
|
adaptive_steps.change( |
|
|
fn=lambda value: (gr.Slider(interactive=not value), gr.Slider(interactive=value)), |
|
|
inputs=[adaptive_steps], |
|
|
outputs=[gamma, gamma_offset], |
|
|
).then( |
|
|
fn=lambda value: on_change_update_setting(consts.GE_USE_ADAPTIVE_STEPS, value), |
|
|
inputs=[adaptive_steps], |
|
|
) |
|
|
|
|
|
use_timestep_adaptive_gamma.change( |
|
|
fn=lambda value: on_change_update_setting(consts.GE_USE_TIMESTEP_ADAPTIVE_GAMMA, value), |
|
|
inputs=[use_timestep_adaptive_gamma], |
|
|
) |
|
|
|
|
|
validate_schedule = gr.Checkbox( |
|
|
label="Validate Schedule", |
|
|
value=from_setting_or_default(consts.GE_VALIDATE_SCHEDULE, False), |
|
|
info="Validate the noise schedule (For debugging purposes).", |
|
|
) |
|
|
|
|
|
with gr.Accordion(label="Extended Reverse SDE", open=False): |
|
|
gr.Markdown("Extended reverse SDE sampler.") |
|
|
gr.Markdown("Max stage for extended reverse SDE.") |
|
|
max_stage = gr.Slider( |
|
|
minimum=1, |
|
|
maximum=3, |
|
|
step=1, |
|
|
value=from_setting_or_default(consts.ER_MAX_STAGE, 3), |
|
|
label="Max Stage", |
|
|
) |
|
|
max_stage.change(fn=lambda value: on_change_update_setting(consts.MAX_STAGE, value), inputs=[max_stage]) |
|
|
|
|
|
return [ |
|
|
euler_a_end, |
|
|
dpm_2m_end, |
|
|
ancestral_eta, |
|
|
detail_strength, |
|
|
langevin_strength, |
|
|
max_stage, |
|
|
adaptive_steps, |
|
|
use_timestep_adaptive_gamma, |
|
|
gamma, |
|
|
gamma_offset, |
|
|
validate_schedule, |
|
|
] |
|
|
|
|
|
def get_values_and_apply(self, p: StableDiffusionProcessing, values: dict): |
|
|
for key, value in values.items(): |
|
|
value = self.xyz_cache.pop(key, value) |
|
|
setattr(p, key, value) |
|
|
p.extra_generation_params[key] = value |
|
|
|
|
|
def process_batch( |
|
|
self, |
|
|
p: StableDiffusionProcessing, |
|
|
euler_a_end: float, |
|
|
dpm_2m_end: float, |
|
|
ancestral_eta: float, |
|
|
detail_strength: float, |
|
|
langevin_strength: float, |
|
|
max_stage: int, |
|
|
use_adaptive_steps: bool, |
|
|
use_timestep_adaptive_gamma: bool, |
|
|
gamma: float, |
|
|
gamma_offset: float, |
|
|
validate_schedule: bool, |
|
|
batch_number: int, |
|
|
prompts: list[str], |
|
|
seeds: list[int], |
|
|
subseeds: list[int], |
|
|
): |
|
|
if p.sampler_name == "Adaptive Progressive": |
|
|
self.get_values_and_apply( |
|
|
p, |
|
|
{ |
|
|
consts.AP_EULER_A_END: euler_a_end, |
|
|
consts.AP_DPM_2M_END: dpm_2m_end, |
|
|
consts.AP_ANCESTRAL_ETA: ancestral_eta, |
|
|
consts.AP_DETAIL_STRENGTH: detail_strength, |
|
|
}, |
|
|
) |
|
|
elif p.sampler_name == "Langevin Euler": |
|
|
self.get_values_and_apply(p, {consts.LANGEVIN_STRENGTH: langevin_strength}) |
|
|
elif p.sampler_name == "Gradient Estimation": |
|
|
self.get_values_and_apply( |
|
|
p, |
|
|
{ |
|
|
consts.GE_GAMMA: gamma, |
|
|
consts.GE_GAMMA_OFFSET: gamma_offset, |
|
|
consts.GE_USE_ADAPTIVE_STEPS: use_adaptive_steps, |
|
|
consts.GE_USE_TIMESTEP_ADAPTIVE_GAMMA: use_timestep_adaptive_gamma, |
|
|
consts.GE_VALIDATE_SCHEDULE: validate_schedule, |
|
|
}, |
|
|
) |
|
|
elif p.sampler_name == "Extended Reverse SDE": |
|
|
self.get_values_and_apply(p, {consts.ER_MAX_STAGE: max_stage}) |
|
|
|
|
|
|
|
|
section = ("exs", "Extra Samplers") |
|
|
|
|
|
|
|
|
def on_settings(): |
|
|
opts.add_option( |
|
|
consts.AP_EULER_A_END, |
|
|
OptionInfo( |
|
|
0.35, |
|
|
"Euler A End", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 0.0, "maximum": 1.0, "step": 0.05}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
opts.add_option( |
|
|
consts.AP_DPM_2M_END, |
|
|
OptionInfo( |
|
|
0.75, |
|
|
"DPM++ 2M End", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 0.0, "maximum": 1.0, "step": 0.05}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
opts.add_option( |
|
|
consts.AP_ANCESTRAL_ETA, |
|
|
OptionInfo( |
|
|
0.4, |
|
|
"Adaptive Progressive Eta", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 0.0, "maximum": 1.0, "step": 0.01}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
opts.add_option( |
|
|
consts.AP_DETAIL_STRENGTH, |
|
|
OptionInfo( |
|
|
1.5, |
|
|
"Adaptive Progressive Detail Strength", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 0.0, "maximum": 3.0, "step": 0.01}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.LANGEVIN_STRENGTH, |
|
|
OptionInfo( |
|
|
0.1, |
|
|
"Langevin Strength", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 0.0, "maximum": 1.0, "step": 0.01}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.ER_MAX_STAGE, |
|
|
OptionInfo( |
|
|
3, |
|
|
"Extended Reverse Time Max Stage", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": 1, "maximum": 3, "step": 1}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.GE_GAMMA, |
|
|
OptionInfo( |
|
|
consts.GE_DEFAULT_GAMMA, |
|
|
"Gradient Estimation Gamma", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": consts.GE_MIN_GAMMA, "maximum": consts.GE_MAX_GAMMA, "step": 0.1}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.GE_GAMMA_OFFSET, |
|
|
OptionInfo( |
|
|
consts.GE_DEFAULT_GAMMA_OFFSET, |
|
|
"Gradient Estimation Gamma Offset", |
|
|
component=gr.Slider, |
|
|
component_args={"minimum": consts.GE_MIN_GAMMA_OFFSET, "maximum": consts.GE_MAX_GAMMA_OFFSET, "step": 0.1}, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.GE_USE_ADAPTIVE_STEPS, |
|
|
OptionInfo( |
|
|
False, |
|
|
"Use Adaptive Steps", |
|
|
component=gr.Checkbox, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
opts.add_option( |
|
|
consts.GE_USE_TIMESTEP_ADAPTIVE_GAMMA, |
|
|
OptionInfo( |
|
|
False, |
|
|
"Use Timestep Adaptive Gamma", |
|
|
component=gr.Checkbox, |
|
|
section=section, |
|
|
), |
|
|
) |
|
|
|
|
|
|
|
|
on_ui_settings(on_settings) |
|
|
|
