__init__.py 8.28 KB
Newer Older
1
__version__ = "0.18.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_k_diffusion_available,
9
    is_k_diffusion_version,
10
    is_librosa_available,
11
    is_note_seq_available,
12
13
    is_onnx_available,
    is_scipy_available,
14
    is_torch_available,
15
    is_torchsde_available,
16
    is_transformers_available,
17
    is_transformers_version,
18
    is_unidecode_available,
19
    logging,
20
)
Patrick von Platen's avatar
Patrick von Platen committed
21
22


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

31
32
33
34
35
36
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
37
38
    from .models import (
        AutoencoderKL,
39
        ControlNetModel,
40
        ModelMixin,
Will Berman's avatar
Will Berman committed
41
        PriorTransformer,
42
        T5FilmDecoder,
Will Berman's avatar
Will Berman committed
43
44
45
46
        Transformer2DModel,
        UNet1DModel,
        UNet2DConditionModel,
        UNet2DModel,
47
        UNet3DConditionModel,
Will Berman's avatar
Will Berman committed
48
49
        VQModel,
    )
50
51
52
53
54
55
56
57
58
    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,
    )
59
    from .pipelines import (
60
        AudioPipelineOutput,
61
62
63
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
64
        DiffusionPipeline,
Kashif Rasul's avatar
Kashif Rasul committed
65
        DiTPipeline,
66
        ImagePipelineOutput,
67
68
        KarrasVePipeline,
        LDMPipeline,
69
        LDMSuperResolutionPipeline,
70
        PNDMPipeline,
Revist's avatar
Revist committed
71
        RePaintPipeline,
72
73
        ScoreSdeVePipeline,
    )
74
    from .schedulers import (
75
        DDIMInverseScheduler,
76
        DDIMParallelScheduler,
77
        DDIMScheduler,
78
        DDPMParallelScheduler,
79
        DDPMScheduler,
80
        DEISMultistepScheduler,
81
        DPMSolverMultistepInverseScheduler,
82
        DPMSolverMultistepScheduler,
83
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
84
85
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
86
        HeunDiscreteScheduler,
87
        IPNDMScheduler,
88
        KarrasVeScheduler,
89
90
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
91
        PNDMScheduler,
Revist's avatar
Revist committed
92
        RePaintScheduler,
93
94
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
95
        UnCLIPScheduler,
96
        UniPCMultistepScheduler,
Will Berman's avatar
Will Berman committed
97
        VQDiffusionScheduler,
98
99
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
100

101
102
103
104
try:
    if not (is_torch_available() and is_scipy_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
105
    from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
106
107
else:
    from .schedulers import LMSDiscreteScheduler
Patrick von Platen's avatar
Patrick von Platen committed
108

109
110
111
112
113
114
115
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
116
117
118
119
120
121
122

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:
123
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
124
125
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
126
        AudioLDMPipeline,
127
        CycleDiffusionPipeline,
Patrick von Platen's avatar
Patrick von Platen committed
128
129
130
131
132
133
        IFImg2ImgPipeline,
        IFImg2ImgSuperResolutionPipeline,
        IFInpaintingPipeline,
        IFInpaintingSuperResolutionPipeline,
        IFPipeline,
        IFSuperResolutionPipeline,
134
        ImageTextPipelineOutput,
YiYi Xu's avatar
YiYi Xu committed
135
136
137
138
        KandinskyImg2ImgPipeline,
        KandinskyInpaintPipeline,
        KandinskyPipeline,
        KandinskyPriorPipeline,
139
        LDMTextToImagePipeline,
140
        PaintByExamplePipeline,
141
        SemanticStableDiffusionPipeline,
YiYi Xu's avatar
YiYi Xu committed
142
        StableDiffusionAttendAndExcitePipeline,
143
144
        StableDiffusionControlNetImg2ImgPipeline,
        StableDiffusionControlNetInpaintPipeline,
145
        StableDiffusionControlNetPipeline,
146
        StableDiffusionDepth2ImgPipeline,
147
        StableDiffusionDiffEditPipeline,
148
        StableDiffusionImageVariationPipeline,
149
150
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
151
        StableDiffusionInpaintPipelineLegacy,
152
        StableDiffusionInstructPix2PixPipeline,
153
        StableDiffusionLatentUpscalePipeline,
estelleafl's avatar
estelleafl committed
154
        StableDiffusionLDM3DPipeline,
155
        StableDiffusionModelEditingPipeline,
156
        StableDiffusionPanoramaPipeline,
157
        StableDiffusionParadigmsPipeline,
158
        StableDiffusionPipeline,
159
        StableDiffusionPipelineSafe,
160
        StableDiffusionPix2PixZeroPipeline,
161
        StableDiffusionSAGPipeline,
162
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
163
164
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
165
        TextToVideoSDPipeline,
166
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
167
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
168
        UnCLIPPipeline,
169
170
171
        UniDiffuserModel,
        UniDiffuserPipeline,
        UniDiffuserTextDecoder,
172
173
174
175
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
176
        VQDiffusionPipeline,
177
    )
178

179
180
181
182
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
183
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
184
185
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
186

187
188
189
190
191
192
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:
193
194
195
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
196
        OnnxStableDiffusionInpaintPipelineLegacy,
197
        OnnxStableDiffusionPipeline,
198
        OnnxStableDiffusionUpscalePipeline,
199
200
        StableDiffusionOnnxPipeline,
    )
201
202
203
204
205
206

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

210
try:
211
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
212
213
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
214
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
215
216
217
else:
    from .pipelines import SpectrogramDiffusionPipeline

218
219
220
221
222
223
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
224
    from .models.controlnet_flax import FlaxControlNetModel
225
    from .models.modeling_flax_utils import FlaxModelMixin
226
227
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
228
    from .pipelines import FlaxDiffusionPipeline
229
230
231
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
232
        FlaxDPMSolverMultistepScheduler,
233
234
235
236
237
238
239
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

240

241
242
243
244
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
245
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
246
else:
247
    from .pipelines import (
YiYi Xu's avatar
YiYi Xu committed
248
        FlaxStableDiffusionControlNetPipeline,
249
250
251
252
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )
253
254
255
256
257
258
259
260

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