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

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

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

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

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

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

211

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

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