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

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

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

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

164
165
166
167
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
168
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
169
170
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
171

172
173
174
175
176
177
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:
178
179
180
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
181
        OnnxStableDiffusionInpaintPipelineLegacy,
182
        OnnxStableDiffusionPipeline,
183
        OnnxStableDiffusionUpscalePipeline,
184
185
        StableDiffusionOnnxPipeline,
    )
186
187
188
189
190
191

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

195
try:
196
    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
197
198
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
199
    from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
200
201
202
else:
    from .pipelines import SpectrogramDiffusionPipeline

203
204
205
206
207
208
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
209
    from .models.controlnet_flax import FlaxControlNetModel
210
    from .models.modeling_flax_utils import FlaxModelMixin
211
212
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
213
    from .pipelines import FlaxDiffusionPipeline
214
215
216
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
217
        FlaxDPMSolverMultistepScheduler,
218
219
220
221
222
223
224
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

225

226
227
228
229
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
230
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
231
else:
232
    from .pipelines import (
YiYi Xu's avatar
YiYi Xu committed
233
        FlaxStableDiffusionControlNetPipeline,
234
235
236
237
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )
238
239
240
241
242
243
244
245

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