__init__.py 6.31 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
12
    is_onnx_available,
    is_scipy_available,
13
    is_torch_available,
14
    is_transformers_available,
15
    is_transformers_version,
16
    is_unidecode_available,
17
    logging,
18
)
Patrick von Platen's avatar
Patrick von Platen committed
19
20


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

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

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

102
103
104
105
106
107
108

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

142
143
144
145
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
146
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
147
148
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
149

150
151
152
153
154
155
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:
156
157
158
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
159
        OnnxStableDiffusionInpaintPipelineLegacy,
160
161
162
        OnnxStableDiffusionPipeline,
        StableDiffusionOnnxPipeline,
    )
163
164
165
166
167
168

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

172
173
174
175
176
177
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
178
    from .models.modeling_flax_utils import FlaxModelMixin
179
180
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
181
    from .pipelines import FlaxDiffusionPipeline
182
183
184
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
185
        FlaxDPMSolverMultistepScheduler,
186
187
188
189
190
191
192
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

193

194
195
196
197
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
198
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
199
else:
200
201
202
203
204
    from .pipelines import (
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )