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] ) # Update interactivity when adaptive steps checkbox changes 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)