__init__.py 9.59 KB
Newer Older
Sayak Paul's avatar
Sayak Paul committed
1
__version__ = "0.21.0.dev0"
2
3

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


24
25
26
27
28
29
30
31
try:
    if not is_onnx_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_onnx_objects import *  # noqa F403
else:
    from .pipelines import OnnxRuntimeModel

32
33
34
35
36
37
try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_pt_objects import *  # noqa F403
else:
Will Berman's avatar
Will Berman committed
38
    from .models import (
Ruslan Vorovchenko's avatar
Ruslan Vorovchenko committed
39
        AsymmetricAutoencoderKL,
Will Berman's avatar
Will Berman committed
40
        AutoencoderKL,
41
        AutoencoderTiny,
42
        ControlNetModel,
43
        ModelMixin,
Will Berman's avatar
Will Berman committed
44
        MultiAdapter,
Will Berman's avatar
Will Berman committed
45
        PriorTransformer,
Will Berman's avatar
Will Berman committed
46
        T2IAdapter,
47
        T5FilmDecoder,
Will Berman's avatar
Will Berman committed
48
49
50
51
        Transformer2DModel,
        UNet1DModel,
        UNet2DConditionModel,
        UNet2DModel,
52
        UNet3DConditionModel,
Will Berman's avatar
Will Berman committed
53
54
        VQModel,
    )
55
56
57
58
59
60
61
62
63
    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,
    )
64
    from .pipelines import (
65
        AudioPipelineOutput,
YiYi Xu's avatar
YiYi Xu committed
66
67
68
        AutoPipelineForImage2Image,
        AutoPipelineForInpainting,
        AutoPipelineForText2Image,
69
        ConsistencyModelPipeline,
70
71
72
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
73
        DiffusionPipeline,
Kashif Rasul's avatar
Kashif Rasul committed
74
        DiTPipeline,
75
        ImagePipelineOutput,
76
77
        KarrasVePipeline,
        LDMPipeline,
78
        LDMSuperResolutionPipeline,
79
        PNDMPipeline,
Revist's avatar
Revist committed
80
        RePaintPipeline,
81
82
        ScoreSdeVePipeline,
    )
83
    from .schedulers import (
84
        CMStochasticIterativeScheduler,
85
        DDIMInverseScheduler,
86
        DDIMParallelScheduler,
87
        DDIMScheduler,
88
        DDPMParallelScheduler,
89
        DDPMScheduler,
90
        DEISMultistepScheduler,
91
        DPMSolverMultistepInverseScheduler,
92
        DPMSolverMultistepScheduler,
93
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
94
95
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
96
        HeunDiscreteScheduler,
97
        IPNDMScheduler,
98
        KarrasVeScheduler,
99
100
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
101
        PNDMScheduler,
Revist's avatar
Revist committed
102
        RePaintScheduler,
103
104
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
105
        UnCLIPScheduler,
106
        UniPCMultistepScheduler,
Will Berman's avatar
Will Berman committed
107
        VQDiffusionScheduler,
108
109
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
110

111
112
113
114
try:
    if not (is_torch_available() and is_scipy_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
115
    from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
116
117
else:
    from .schedulers import LMSDiscreteScheduler
Patrick von Platen's avatar
Patrick von Platen committed
118

119
120
121
122
123
124
125
try:
    if not (is_torch_available() and is_torchsde_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
    from .schedulers import DPMSolverSDEScheduler
126
127
128
129
130
131
132

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:
133
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
134
135
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
136
137
138
        AudioLDM2Pipeline,
        AudioLDM2ProjectionModel,
        AudioLDM2UNet2DConditionModel,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
139
        AudioLDMPipeline,
140
        CycleDiffusionPipeline,
Patrick von Platen's avatar
Patrick von Platen committed
141
142
143
144
145
146
        IFImg2ImgPipeline,
        IFImg2ImgSuperResolutionPipeline,
        IFInpaintingPipeline,
        IFInpaintingSuperResolutionPipeline,
        IFPipeline,
        IFSuperResolutionPipeline,
147
        ImageTextPipelineOutput,
148
149
        KandinskyCombinedPipeline,
        KandinskyImg2ImgCombinedPipeline,
YiYi Xu's avatar
YiYi Xu committed
150
        KandinskyImg2ImgPipeline,
151
        KandinskyInpaintCombinedPipeline,
YiYi Xu's avatar
YiYi Xu committed
152
153
154
        KandinskyInpaintPipeline,
        KandinskyPipeline,
        KandinskyPriorPipeline,
155
        KandinskyV22CombinedPipeline,
YiYi Xu's avatar
YiYi Xu committed
156
157
        KandinskyV22ControlnetImg2ImgPipeline,
        KandinskyV22ControlnetPipeline,
158
        KandinskyV22Img2ImgCombinedPipeline,
YiYi Xu's avatar
YiYi Xu committed
159
        KandinskyV22Img2ImgPipeline,
160
        KandinskyV22InpaintCombinedPipeline,
YiYi Xu's avatar
YiYi Xu committed
161
162
163
164
        KandinskyV22InpaintPipeline,
        KandinskyV22Pipeline,
        KandinskyV22PriorEmb2EmbPipeline,
        KandinskyV22PriorPipeline,
165
        LDMTextToImagePipeline,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
166
        MusicLDMPipeline,
167
        PaintByExamplePipeline,
168
        SemanticStableDiffusionPipeline,
YiYi Xu's avatar
YiYi Xu committed
169
170
        ShapEImg2ImgPipeline,
        ShapEPipeline,
Will Berman's avatar
Will Berman committed
171
        StableDiffusionAdapterPipeline,
YiYi Xu's avatar
YiYi Xu committed
172
        StableDiffusionAttendAndExcitePipeline,
173
174
        StableDiffusionControlNetImg2ImgPipeline,
        StableDiffusionControlNetInpaintPipeline,
175
        StableDiffusionControlNetPipeline,
176
        StableDiffusionDepth2ImgPipeline,
177
        StableDiffusionDiffEditPipeline,
178
        StableDiffusionGLIGENPipeline,
179
        StableDiffusionImageVariationPipeline,
180
181
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
182
        StableDiffusionInpaintPipelineLegacy,
183
        StableDiffusionInstructPix2PixPipeline,
184
        StableDiffusionLatentUpscalePipeline,
estelleafl's avatar
estelleafl committed
185
        StableDiffusionLDM3DPipeline,
186
        StableDiffusionModelEditingPipeline,
187
        StableDiffusionPanoramaPipeline,
188
        StableDiffusionParadigmsPipeline,
189
        StableDiffusionPipeline,
190
        StableDiffusionPipelineSafe,
191
        StableDiffusionPix2PixZeroPipeline,
192
        StableDiffusionSAGPipeline,
193
        StableDiffusionUpscalePipeline,
194
195
196
197
198
        StableDiffusionXLControlNetPipeline,
        StableDiffusionXLImg2ImgPipeline,
        StableDiffusionXLInpaintPipeline,
        StableDiffusionXLInstructPix2PixPipeline,
        StableDiffusionXLPipeline,
Will Berman's avatar
Will Berman committed
199
200
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
201
        TextToVideoSDPipeline,
202
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
203
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
204
        UnCLIPPipeline,
205
206
207
        UniDiffuserModel,
        UniDiffuserPipeline,
        UniDiffuserTextDecoder,
208
209
210
211
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
Patrick von Platen's avatar
Patrick von Platen committed
212
        VideoToVideoSDPipeline,
213
        VQDiffusionPipeline,
214
    )
215

216
217
218
219
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
220
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
221
222
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
223

224
225
226
227
228
229
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:
230
231
232
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
233
        OnnxStableDiffusionInpaintPipelineLegacy,
234
        OnnxStableDiffusionPipeline,
235
        OnnxStableDiffusionUpscalePipeline,
236
237
        StableDiffusionOnnxPipeline,
    )
238
239
240
241
242
243

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

247
try:
248
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
249
250
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
251
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
252
253
254
else:
    from .pipelines import SpectrogramDiffusionPipeline

255
256
257
258
259
260
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
YiYi Xu's avatar
YiYi Xu committed
261
    from .models.controlnet_flax import FlaxControlNetModel
262
    from .models.modeling_flax_utils import FlaxModelMixin
263
264
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
265
    from .pipelines import FlaxDiffusionPipeline
266
267
268
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
269
        FlaxDPMSolverMultistepScheduler,
270
271
272
273
274
275
276
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

277

278
279
280
281
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
282
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
283
else:
284
    from .pipelines import (
YiYi Xu's avatar
YiYi Xu committed
285
        FlaxStableDiffusionControlNetPipeline,
286
287
288
289
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )
290
291
292
293
294
295
296
297

try:
    if not (is_note_seq_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_note_seq_objects import *  # noqa F403
else:
    from .pipelines import MidiProcessor