__init__.py 7.14 KB
Newer Older
1
__version__ = "0.15.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_transformers_available,
16
    is_transformers_version,
17
    is_unidecode_available,
18
    logging,
19
)
Patrick von Platen's avatar
Patrick von Platen committed
20
21


22
23
24
25
26
27
28
29
try:
    if not is_onnx_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_onnx_objects import *  # noqa F403
else:
    from .pipelines import OnnxRuntimeModel

30
31
32
33
34
35
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
36
37
    from .models import (
        AutoencoderKL,
38
        ControlNetModel,
39
        ModelMixin,
Will Berman's avatar
Will Berman committed
40
        PriorTransformer,
41
        T5FilmDecoder,
Will Berman's avatar
Will Berman committed
42
43
44
45
        Transformer2DModel,
        UNet1DModel,
        UNet2DConditionModel,
        UNet2DModel,
46
        UNet3DConditionModel,
Will Berman's avatar
Will Berman committed
47
48
        VQModel,
    )
49
50
51
52
53
54
55
56
57
    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,
    )
58
    from .pipelines import (
59
        AudioPipelineOutput,
60
61
62
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
63
        DiffusionPipeline,
Kashif Rasul's avatar
Kashif Rasul committed
64
        DiTPipeline,
65
        ImagePipelineOutput,
66
67
        KarrasVePipeline,
        LDMPipeline,
68
        LDMSuperResolutionPipeline,
69
        PNDMPipeline,
Revist's avatar
Revist committed
70
        RePaintPipeline,
71
72
        ScoreSdeVePipeline,
    )
73
    from .schedulers import (
74
        DDIMInverseScheduler,
75
76
        DDIMScheduler,
        DDPMScheduler,
77
        DEISMultistepScheduler,
78
        DPMSolverMultistepScheduler,
79
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
80
81
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
82
        HeunDiscreteScheduler,
83
        IPNDMScheduler,
84
        KarrasVeScheduler,
85
86
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
87
        PNDMScheduler,
Revist's avatar
Revist committed
88
        RePaintScheduler,
89
90
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
91
        UnCLIPScheduler,
92
        UniPCMultistepScheduler,
Will Berman's avatar
Will Berman committed
93
        VQDiffusionScheduler,
94
95
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
96

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

105
106
107
108
109
110
111

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:
112
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
113
114
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
Sanchit Gandhi's avatar
Sanchit Gandhi committed
115
        AudioLDMPipeline,
116
        CycleDiffusionPipeline,
117
        LDMTextToImagePipeline,
118
        PaintByExamplePipeline,
119
        SemanticStableDiffusionPipeline,
YiYi Xu's avatar
YiYi Xu committed
120
        StableDiffusionAttendAndExcitePipeline,
121
        StableDiffusionControlNetPipeline,
122
        StableDiffusionDepth2ImgPipeline,
123
        StableDiffusionImageVariationPipeline,
124
125
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
126
        StableDiffusionInpaintPipelineLegacy,
127
        StableDiffusionInstructPix2PixPipeline,
128
        StableDiffusionLatentUpscalePipeline,
129
        StableDiffusionModelEditingPipeline,
130
        StableDiffusionPanoramaPipeline,
131
        StableDiffusionPipeline,
132
        StableDiffusionPipelineSafe,
133
        StableDiffusionPix2PixZeroPipeline,
134
        StableDiffusionSAGPipeline,
135
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
136
137
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
138
        TextToVideoSDPipeline,
Will Berman's avatar
Will Berman committed
139
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
140
        UnCLIPPipeline,
141
142
143
144
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
145
        VQDiffusionPipeline,
146
    )
147

148
149
150
151
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
152
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
153
154
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
155

156
157
158
159
160
161
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:
162
163
164
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
165
        OnnxStableDiffusionInpaintPipelineLegacy,
166
        OnnxStableDiffusionPipeline,
167
        OnnxStableDiffusionUpscalePipeline,
168
169
        StableDiffusionOnnxPipeline,
    )
170
171
172
173
174
175

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

179
180
181
182
183
184
185
186
try:
    if not (is_torch_available() and is_note_seq_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_torch_and_note_seq_objects import *  # noqa F403
else:
    from .pipelines import SpectrogramDiffusionPipeline

187
188
189
190
191
192
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
193
    from .models.controlnet_flax import FlaxControlNetModel
194
    from .models.modeling_flax_utils import FlaxModelMixin
195
196
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
197
    from .pipelines import FlaxDiffusionPipeline
198
199
200
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
201
        FlaxDPMSolverMultistepScheduler,
202
203
204
205
206
207
208
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

209

210
211
212
213
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
214
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
215
else:
216
    from .pipelines import (
YiYi Xu's avatar
YiYi Xu committed
217
        FlaxStableDiffusionControlNetPipeline,
218
219
220
221
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )
222
223
224
225
226
227
228
229

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