__init__.py 8.02 KB
Newer Older
1
__version__ = "0.17.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,
YiYi Xu's avatar
YiYi Xu committed
132
133
134
135
        KandinskyImg2ImgPipeline,
        KandinskyInpaintPipeline,
        KandinskyPipeline,
        KandinskyPriorPipeline,
136
        LDMTextToImagePipeline,
137
        PaintByExamplePipeline,
138
        SemanticStableDiffusionPipeline,
YiYi Xu's avatar
YiYi Xu committed
139
        StableDiffusionAttendAndExcitePipeline,
140
141
        StableDiffusionControlNetImg2ImgPipeline,
        StableDiffusionControlNetInpaintPipeline,
142
        StableDiffusionControlNetPipeline,
143
        StableDiffusionDepth2ImgPipeline,
144
        StableDiffusionDiffEditPipeline,
145
        StableDiffusionImageVariationPipeline,
146
147
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
148
        StableDiffusionInpaintPipelineLegacy,
149
        StableDiffusionInstructPix2PixPipeline,
150
        StableDiffusionLatentUpscalePipeline,
151
        StableDiffusionModelEditingPipeline,
152
        StableDiffusionPanoramaPipeline,
153
        StableDiffusionPipeline,
154
        StableDiffusionPipelineSafe,
155
        StableDiffusionPix2PixZeroPipeline,
156
        StableDiffusionSAGPipeline,
157
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
158
159
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
160
        TextToVideoSDPipeline,
161
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
162
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
163
        UnCLIPPipeline,
164
165
166
167
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
168
        VQDiffusionPipeline,
169
    )
170

171
172
173
174
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
175
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
176
177
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
178

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

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

202
try:
203
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
204
205
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
206
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
207
208
209
else:
    from .pipelines import SpectrogramDiffusionPipeline

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

232

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

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