__init__.py 7.22 KB
Newer Older
1
__version__ = "0.16.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,
139
        TextToVideoZeroPipeline,
Will Berman's avatar
Will Berman committed
140
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
141
        UnCLIPPipeline,
142
143
144
145
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
146
        VQDiffusionPipeline,
147
    )
148

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

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

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

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

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

210

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

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