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

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

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

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

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

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

224

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

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