Skip to content

Commit 570d3f1

Browse files
osansevieroanton-l
andauthored
Expose LR schedulers (#80)
* Expose schedulers * Update __init__.py Co-authored-by: Anton Lozhkov <[email protected]>
1 parent 85244d4 commit 570d3f1

File tree

1 file changed

+10
-0
lines changed

1 file changed

+10
-0
lines changed

src/diffusers/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,19 @@
88

99
from .modeling_utils import ModelMixin
1010
from .models import AutoencoderKL, UNet2DConditionModel, UNet2DModel, VQModel
11+
from .optimization import (
12+
get_constant_schedule,
13+
get_constant_schedule_with_warmup,
14+
get_linear_schedule_with_warmup,
15+
get_cosine_schedule_with_warmup,
16+
get_cosine_with_hard_restarts_schedule_with_warmup,
17+
get_polynomial_decay_schedule_with_warmup,
18+
get_scheduler,
19+
)
1120
from .pipeline_utils import DiffusionPipeline
1221
from .pipelines import DDIMPipeline, DDPMPipeline, LDMPipeline, PNDMPipeline, ScoreSdeVePipeline
1322
from .schedulers import DDIMScheduler, DDPMScheduler, PNDMScheduler, SchedulerMixin, ScoreSdeVeScheduler
23+
from .training_utils import EMAModel
1424

1525

1626
if is_transformers_available():

0 commit comments

Comments
 (0)