__init__.py 5.68 KB
Newer Older
anton-'s avatar
anton- committed
1
__version__ = "0.10.0"
2
3
4

from .configuration_utils import ConfigMixin
from .onnx_utils import OnnxRuntimeModel
5
from .utils import (
6
    OptionalDependencyNotAvailable,
7
    is_flax_available,
8
    is_inflect_available,
9
    is_k_diffusion_available,
10
    is_librosa_available,
11
12
    is_onnx_available,
    is_scipy_available,
13
    is_torch_available,
14
    is_transformers_available,
15
    is_transformers_version,
16
    is_unidecode_available,
17
    logging,
18
)
Patrick von Platen's avatar
Patrick von Platen committed
19
20


21
22
23
24
25
26
27
28
29
30
31
32
# Make sure `transformers` is up to date
if is_transformers_available():
    import transformers

    if is_transformers_version("<", "4.25.1"):
        raise ImportError(
            f"`diffusers` requires transformers >= 4.25.1 to function correctly, but {transformers.__version__} was"
            " found in your environment. You can upgrade it with pip: `pip install transformers --upgrade`"
        )
else:
    pass

33
34
35
36
37
38
try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_pt_objects import *  # noqa F403
else:
39
    from .modeling_utils import ModelMixin
Will Berman's avatar
Will Berman committed
40
    from .models import AutoencoderKL, Transformer2DModel, UNet1DModel, UNet2DConditionModel, UNet2DModel, VQModel
41
42
43
44
45
46
47
48
49
50
    from .optimization import (
        get_constant_schedule,
        get_constant_schedule_with_warmup,
        get_cosine_schedule_with_warmup,
        get_cosine_with_hard_restarts_schedule_with_warmup,
        get_linear_schedule_with_warmup,
        get_polynomial_decay_schedule_with_warmup,
        get_scheduler,
    )
    from .pipeline_utils import DiffusionPipeline
51
52
53
54
55
56
    from .pipelines import (
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
        KarrasVePipeline,
        LDMPipeline,
57
        LDMSuperResolutionPipeline,
58
        PNDMPipeline,
Revist's avatar
Revist committed
59
        RePaintPipeline,
60
61
        ScoreSdeVePipeline,
    )
62
63
64
    from .schedulers import (
        DDIMScheduler,
        DDPMScheduler,
65
        DPMSolverMultistepScheduler,
66
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
67
68
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
69
        HeunDiscreteScheduler,
70
        IPNDMScheduler,
71
        KarrasVeScheduler,
72
73
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
74
        PNDMScheduler,
Revist's avatar
Revist committed
75
        RePaintScheduler,
76
77
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
78
        VQDiffusionScheduler,
79
80
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
81

82
83
84
85
try:
    if not (is_torch_available() and is_scipy_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
86
    from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
87
88
else:
    from .schedulers import LMSDiscreteScheduler
Patrick von Platen's avatar
Patrick von Platen committed
89

90
91
92
93
94
95
96

try:
    if not (is_torch_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
97
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
98
99
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
100
        CycleDiffusionPipeline,
101
        LDMTextToImagePipeline,
102
        PaintByExamplePipeline,
103
        StableDiffusionDepth2ImgPipeline,
104
        StableDiffusionImageVariationPipeline,
105
106
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
107
        StableDiffusionInpaintPipelineLegacy,
108
        StableDiffusionPipeline,
109
        StableDiffusionPipelineSafe,
110
        StableDiffusionUpscalePipeline,
111
112
113
114
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
115
        VQDiffusionPipeline,
116
    )
117

118
119
120
121
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
122
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
123
124
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
125

126
127
128
129
130
131
try:
    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
132
133
134
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
135
        OnnxStableDiffusionInpaintPipelineLegacy,
136
137
138
        OnnxStableDiffusionPipeline,
        StableDiffusionOnnxPipeline,
    )
139
140
141
142
143
144

try:
    if not (is_torch_available() and is_librosa_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
145
else:
146
    from .pipelines import AudioDiffusionPipeline, Mel
147

148
149
150
151
152
153
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
154
155
156
157
158
159
160
    from .modeling_flax_utils import FlaxModelMixin
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
    from .pipeline_flax_utils import FlaxDiffusionPipeline
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
161
        FlaxDPMSolverMultistepScheduler,
162
163
164
165
166
167
168
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

169
170
171
172
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
173
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
174
175
else:
    from .pipelines import FlaxStableDiffusionPipeline