Expose LR schedulers (huggingface#80)

* Expose schedulers * Update __init__.py Co-authored-by: Anton Lozhkov <[email protected]>
titu1994 · Jul 22, 2022 · 570d3f1 · 570d3f1
1 parent 85244d4
commit 570d3f1
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
@@ -8,9 +8,19 @@
 
 from .modeling_utils import ModelMixin
 from .models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
+from .optimization import (
+    get_constant_schedule,
+    get_constant_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+    get_cosine_schedule_with_warmup,
+    get_cosine_with_hard_restarts_schedule_with_warmup,
+    get_polynomial_decay_schedule_with_warmup,
+    get_scheduler,
+)
 from .pipeline_utils import DiffusionPipeline
 from .pipelines import DDIMPipeline, DDPMPipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline
 from .schedulers import DDIMScheduler, DDPMScheduler, PNDMScheduler, SchedulerMixin, ScoreSdeVeScheduler
+from .training_utils import EMAModel
 
 
 if is_transformers_available():