Skip to content

Commit a4ec54a

Browse files
Add linear_start and linear_end to model_config.sampling_settings
1 parent 9ac0b48 commit a4ec54a

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

comfy/model_sampling.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,17 @@ def calculate_denoised(self, sigma, model_output, model_input):
2222
class ModelSamplingDiscrete(torch.nn.Module):
2323
def __init__(self, model_config=None):
2424
super().__init__()
25-
beta_schedule = "linear"
25+
2626
if model_config is not None:
27-
beta_schedule = model_config.sampling_settings.get("beta_schedule", beta_schedule)
28-
self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=0.00085, linear_end=0.012, cosine_s=8e-3)
27+
sampling_settings = model_config.sampling_settings
28+
else:
29+
sampling_settings = {}
30+
31+
beta_schedule = sampling_settings.get("beta_schedule", "linear")
32+
linear_start = sampling_settings.get("linear_start", 0.00085)
33+
linear_end = sampling_settings.get("linear_end", 0.012)
34+
35+
self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3)
2936
self.sigma_data = 1.0
3037

3138
def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,

0 commit comments

Comments
 (0)