From 821f93872ecf3f1769b8abebd640a804839f4ab9 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 16 Jul 2024 15:18:24 -0400 Subject: [PATCH] Allow model sampling to set number of timesteps. --- comfy/model_sampling.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py index 2d95a83d..25bb7e04 100644 --- a/comfy/model_sampling.py +++ b/comfy/model_sampling.py @@ -59,8 +59,9 @@ class ModelSamplingDiscrete(torch.nn.Module): beta_schedule = sampling_settings.get("beta_schedule", "linear") linear_start = sampling_settings.get("linear_start", 0.00085) linear_end = sampling_settings.get("linear_end", 0.012) + timesteps = sampling_settings.get("timesteps", 1000) - self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3) + self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3) self.sigma_data = 1.0 def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,