From 570d3f1eb92c723653b9f07177fb10f4bbc15a8b Mon Sep 17 00:00:00 2001 From: Omar Sanseviero Date: Fri, 22 Jul 2022 13:29:42 +0200 Subject: [PATCH] Expose LR schedulers (#80) * Expose schedulers * Update __init__.py Co-authored-by: Anton Lozhkov --- src/diffusers/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 3dc61a13..b2513899 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -8,9 +8,19 @@ __version__ = "0.1.2" 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():