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

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

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

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

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

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

228

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

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