__init__.py 9.56 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,
166
        PaintByExamplePipeline,
167
        SemanticStableDiffusionPipeline,
YiYi Xu's avatar
YiYi Xu committed
168
169
        ShapEImg2ImgPipeline,
        ShapEPipeline,
Will Berman's avatar
Will Berman committed
170
        StableDiffusionAdapterPipeline,
YiYi Xu's avatar
YiYi Xu committed
171
        StableDiffusionAttendAndExcitePipeline,
172
173
        StableDiffusionControlNetImg2ImgPipeline,
        StableDiffusionControlNetInpaintPipeline,
174
        StableDiffusionControlNetPipeline,
175
        StableDiffusionDepth2ImgPipeline,
176
        StableDiffusionDiffEditPipeline,
177
        StableDiffusionGLIGENPipeline,
178
        StableDiffusionImageVariationPipeline,
179
180
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
181
        StableDiffusionInpaintPipelineLegacy,
182
        StableDiffusionInstructPix2PixPipeline,
183
        StableDiffusionLatentUpscalePipeline,
estelleafl's avatar
estelleafl committed
184
        StableDiffusionLDM3DPipeline,
185
        StableDiffusionModelEditingPipeline,
186
        StableDiffusionPanoramaPipeline,
187
        StableDiffusionParadigmsPipeline,
188
        StableDiffusionPipeline,
189
        StableDiffusionPipelineSafe,
190
        StableDiffusionPix2PixZeroPipeline,
191
        StableDiffusionSAGPipeline,
192
        StableDiffusionUpscalePipeline,
193
194
195
196
197
        StableDiffusionXLControlNetPipeline,
        StableDiffusionXLImg2ImgPipeline,
        StableDiffusionXLInpaintPipeline,
        StableDiffusionXLInstructPix2PixPipeline,
        StableDiffusionXLPipeline,
Will Berman's avatar
Will Berman committed
198
199
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
200
        TextToVideoSDPipeline,
201
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
202
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
203
        UnCLIPPipeline,
204
205
206
        UniDiffuserModel,
        UniDiffuserPipeline,
        UniDiffuserTextDecoder,
207
208
209
210
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
Patrick von Platen's avatar
Patrick von Platen committed
211
        VideoToVideoSDPipeline,
212
        VQDiffusionPipeline,
213
    )
214

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

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

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

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

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

276

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

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