__init__.py 7.24 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 .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,
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