From 9de7084884f4266a19e4ed3f503687927c845c4d Mon Sep 17 00:00:00 2001
From: v0xie <28695009+v0xie@users.noreply.github.com>
Date: Fri, 19 Jul 2024 14:54:24 -0700
Subject: [PATCH] always add alpha/beta to extra_generation_params when
 schedule is Beta

---
 modules/sd_samplers_kdiffusion.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index 95a354dac..4e5310bc6 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -119,6 +119,10 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
 
             if scheduler.need_inner_model:
                 sigmas_kwargs['inner_model'] = self.model_wrap
+            
+            if scheduler.label == 'Beta':
+                p.extra_generation_params["Beta schedule alpha"] = opts.beta_dist_alpha
+                p.extra_generation_params["Beta schedule beta"] = opts.beta_dist_beta
 
             sigmas = scheduler.function(n=steps, **sigmas_kwargs, device=devices.cpu)