__init__.py 9.94 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
        CLIPImageProjection,
70
        ConsistencyModelPipeline,
71
72
73
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
74
        DiffusionPipeline,
Kashif Rasul's avatar
Kashif Rasul committed
75
        DiTPipeline,
76
        ImagePipelineOutput,
77
78
        KarrasVePipeline,
        LDMPipeline,
79
        LDMSuperResolutionPipeline,
80
        PNDMPipeline,
Revist's avatar
Revist committed
81
        RePaintPipeline,
82
83
        ScoreSdeVePipeline,
    )
84
    from .schedulers import (
85
        CMStochasticIterativeScheduler,
86
        DDIMInverseScheduler,
87
        DDIMParallelScheduler,
88
        DDIMScheduler,
89
        DDPMParallelScheduler,
90
        DDPMScheduler,
Kashif Rasul's avatar
Kashif Rasul committed
91
        DDPMWuerstchenScheduler,
92
        DEISMultistepScheduler,
93
        DPMSolverMultistepInverseScheduler,
94
        DPMSolverMultistepScheduler,
95
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
96
97
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
98
        HeunDiscreteScheduler,
99
        IPNDMScheduler,
100
        KarrasVeScheduler,
101
102
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
103
        PNDMScheduler,
Revist's avatar
Revist committed
104
        RePaintScheduler,
105
106
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
107
        UnCLIPScheduler,
108
        UniPCMultistepScheduler,
Will Berman's avatar
Will Berman committed
109
        VQDiffusionScheduler,
110
111
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
112

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

121
122
123
124
125
126
127
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
128
129
130
131
132
133
134

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

225
226
227
228
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
229
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
230
231
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
232

233
234
235
236
237
238
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:
239
240
241
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
242
        OnnxStableDiffusionInpaintPipelineLegacy,
243
        OnnxStableDiffusionPipeline,
244
        OnnxStableDiffusionUpscalePipeline,
245
246
        StableDiffusionOnnxPipeline,
    )
247
248
249
250
251
252

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

256
try:
257
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
258
259
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
260
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
261
262
263
else:
    from .pipelines import SpectrogramDiffusionPipeline

264
265
266
267
268
269
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
270
    from .models.controlnet_flax import FlaxControlNetModel
271
    from .models.modeling_flax_utils import FlaxModelMixin
272
273
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
274
    from .pipelines import FlaxDiffusionPipeline
275
276
277
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
278
        FlaxDPMSolverMultistepScheduler,
279
280
281
282
283
284
285
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

286

287
288
289
290
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
291
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
292
else:
293
    from .pipelines import (
YiYi Xu's avatar
YiYi Xu committed
294
        FlaxStableDiffusionControlNetPipeline,
295
296
297
298
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )
299
300
301
302
303
304
305
306

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