__init__.py 8.14 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
77
        DDIMScheduler,
        DDPMScheduler,
78
        DEISMultistepScheduler,
79
        DPMSolverMultistepInverseScheduler,
80
        DPMSolverMultistepScheduler,
81
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
82
83
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
84
        HeunDiscreteScheduler,
85
        IPNDMScheduler,
86
        KarrasVeScheduler,
87
88
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
89
        PNDMScheduler,
Revist's avatar
Revist committed
90
        RePaintScheduler,
91
92
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
93
        UnCLIPScheduler,
94
        UniPCMultistepScheduler,
Will Berman's avatar
Will Berman committed
95
        VQDiffusionScheduler,
96
97
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
98

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

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

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

175
176
177
178
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
179
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
180
181
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
182

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

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

206
try:
207
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
208
209
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
210
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
211
212
213
else:
    from .pipelines import SpectrogramDiffusionPipeline

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

236

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

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