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

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

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

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

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

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

227

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

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