diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index cfe71286..7dfb518b 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -200,6 +200,31 @@ def bong_tangent_scheduler(n, sigma_min, sigma_max, device, *, start=1.0, middle return tan_sigmas.to(device) +def flow_match_euler_discrete_scheduler(n, sigma_min, sigma_max, inner_model, device): + from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + unet = inner_model.inner_model.forge_objects.unet + + config = { + "num_train_timesteps": 1000, + "shift": getattr(unet.model.predictor, "shift", 1.0), + "use_dynamic_shifting": shared.opts.use_dynamic_shifting, + "invert_sigmas": shared.opts.invert_sigmas, + "shift_terminal": None, + "use_karras_sigmas": shared.opts.use_karras_sigmas, + "use_exponential_sigmas": shared.opts.use_exponential_sigmas, + "use_beta_sigmas": shared.opts.use_beta_sigmas, + "time_shift_type": "exponential", + "stochastic_sampling": shared.opts.stochastic_sampling, + } + + scheduler = FlowMatchEulerDiscreteScheduler.from_config(config) + scheduler.set_timesteps(n, device=device, mu=0.0) + sigmas = scheduler.sigmas + + return torch.FloatTensor(sigmas).to(device) + + schedulers = [ Scheduler("automatic", "Automatic", None), Scheduler("karras", "Karras", k_diffusion.sampling.get_sigmas_karras, default_rho=7.0), @@ -216,6 +241,7 @@ schedulers = [ Scheduler("beta", "Beta", beta_scheduler, need_inner_model=True), Scheduler("turbo", "Turbo", turbo_scheduler, need_inner_model=True), Scheduler("bong_tangent", "Bong Tangent", bong_tangent_scheduler), + Scheduler("flow_match", "FlowMatchEulerDiscrete", flow_match_euler_discrete_scheduler, need_inner_model=True), ] schedulers_map = {**{x.name: x for x in schedulers}, **{x.label: x for x in schedulers}} diff --git a/modules/shared_options.py b/modules/shared_options.py index 648296d0..d26a8015 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -551,6 +551,12 @@ options_templates.update( "sd_noise_schedule": OptionInfo("Default", "Noise schedule for sampling", gr.Radio, {"choices": ("Default", "Zero Terminal SNR")}, infotext="Noise Schedule"), "beta_dist_alpha": OptionInfo(0.6, "Beta scheduler - alpha", gr.Slider, {"minimum": 0.01, "maximum": 2.0, "step": 0.01}, infotext="Beta scheduler alpha"), "beta_dist_beta": OptionInfo(0.6, "Beta scheduler - beta", gr.Slider, {"minimum": 0.01, "maximum": 2.0, "step": 0.01}, infotext="Beta scheduler beta"), + "use_dynamic_shifting": OptionInfo(False, "use_dynamic_shifting"), + "invert_sigmas": OptionInfo(False, "invert_sigmas"), + "use_karras_sigmas": OptionInfo(False, "use_karras_sigmas"), + "use_exponential_sigmas": OptionInfo(False, "use_exponential_sigmas"), + "use_beta_sigmas": OptionInfo(False, "use_beta_sigmas"), + "stochastic_sampling": OptionInfo(False, "stochastic_sampling"), }, ) )