From 64b1e07e55caeaa95072f6826940c4049687d8bb Mon Sep 17 00:00:00 2001 From: William Berman Date: Wed, 15 Feb 2023 10:27:44 -0800 Subject: [PATCH] KarrasDiffusionSchedulers type note --- src/diffusers/schedulers/scheduling_utils.py | 4 ++++ src/diffusers/schedulers/scheduling_utils_flax.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py index f4103d4d62cc..424447efbdfc 100644 --- a/src/diffusers/schedulers/scheduling_utils.py +++ b/src/diffusers/schedulers/scheduling_utils.py @@ -25,6 +25,10 @@ SCHEDULER_CONFIG_NAME = "scheduler_config.json" +# NOTE: We make this type an enum because it simplifies usage in docs and prevents +# circular imports when used for `_compatibles` within the schedulers module. +# When it's used as a type in pipelines, it really is a Union because the actual +# scheduler instance is passed in. class KarrasDiffusionSchedulers(Enum): DDIMScheduler = 1 DDPMScheduler = 2 diff --git a/src/diffusers/schedulers/scheduling_utils_flax.py b/src/diffusers/schedulers/scheduling_utils_flax.py index 9708c0883760..67c9dcc62414 100644 --- a/src/diffusers/schedulers/scheduling_utils_flax.py +++ b/src/diffusers/schedulers/scheduling_utils_flax.py @@ -27,6 +27,10 @@ SCHEDULER_CONFIG_NAME = "scheduler_config.json" +# NOTE: We make this type an enum because it simplifies usage in docs and prevents +# circular imports when used for `_compatibles` within the schedulers module. +# When it's used as a type in pipelines, it really is a Union because the actual +# scheduler instance is passed in. class FlaxKarrasDiffusionSchedulers(Enum): FlaxDDIMScheduler = 1 FlaxDDPMScheduler = 2