diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py index 37a3ac725..9e2a1c1af 100644 --- a/comfy/model_sampling.py +++ b/comfy/model_sampling.py @@ -24,7 +24,7 @@ class ModelSamplingDiscrete(torch.nn.Module): super().__init__() beta_schedule = "linear" if model_config is not None: - beta_schedule = model_config.beta_schedule + beta_schedule = model_config.sampling_settings.get("beta_schedule", beta_schedule) self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=0.00085, linear_end=0.012, cosine_s=8e-3) self.sigma_data = 1.0 diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py index b073eb4fc..3412cfea0 100644 --- a/comfy/supported_models_base.py +++ b/comfy/supported_models_base.py @@ -19,7 +19,7 @@ class BASE: clip_prefix = [] clip_vision_prefix = None noise_aug_config = None - beta_schedule = "linear" + sampling_settings = {} latent_format = latent_formats.LatentFormat @classmethod