__init__.py 9.64 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
        StableDiffusionXLControlNetImg2ImgPipeline,
195
196
197
198
199
        StableDiffusionXLControlNetPipeline,
        StableDiffusionXLImg2ImgPipeline,
        StableDiffusionXLInpaintPipeline,
        StableDiffusionXLInstructPix2PixPipeline,
        StableDiffusionXLPipeline,
Will Berman's avatar
Will Berman committed
200
201
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
202
        TextToVideoSDPipeline,
203
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
204
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
205
        UnCLIPPipeline,
206
207
208
        UniDiffuserModel,
        UniDiffuserPipeline,
        UniDiffuserTextDecoder,
209
210
211
212
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
Patrick von Platen's avatar
Patrick von Platen committed
213
        VideoToVideoSDPipeline,
214
        VQDiffusionPipeline,
215
    )
216

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

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

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

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

256
257
258
259
260
261
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
262
    from .models.controlnet_flax import FlaxControlNetModel
263
    from .models.modeling_flax_utils import FlaxModelMixin
264
265
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
266
    from .pipelines import FlaxDiffusionPipeline
267
268
269
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
270
        FlaxDPMSolverMultistepScheduler,
271
272
273
274
275
276
277
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

278

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

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