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

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

99
100
101
102
103
104
105

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:
106
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
107
108
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
109
        CycleDiffusionPipeline,
110
        LDMTextToImagePipeline,
111
        PaintByExamplePipeline,
112
        StableDiffusionDepth2ImgPipeline,
113
        StableDiffusionImageVariationPipeline,
114
115
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
116
        StableDiffusionInpaintPipelineLegacy,
117
        StableDiffusionInstructPix2PixPipeline,
118
        StableDiffusionLatentUpscalePipeline,
119
        StableDiffusionPipeline,
120
        StableDiffusionPipelineSafe,
121
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
122
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
123
        UnCLIPPipeline,
124
125
126
127
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
128
        VQDiffusionPipeline,
129
    )
130

131
132
133
134
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
135
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
136
137
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
138

139
140
141
142
143
144
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:
145
146
147
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
148
        OnnxStableDiffusionInpaintPipelineLegacy,
149
150
151
        OnnxStableDiffusionPipeline,
        StableDiffusionOnnxPipeline,
    )
152
153
154
155
156
157

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

161
162
163
164
165
166
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
167
    from .models.modeling_flax_utils import FlaxModelMixin
168
169
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
170
    from .pipelines import FlaxDiffusionPipeline
171
172
173
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
174
        FlaxDPMSolverMultistepScheduler,
175
176
177
178
179
180
181
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

182

183
184
185
186
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
187
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
188
else:
189
190
191
192
193
    from .pipelines import (
        FlaxStableDiffusionImg2ImgPipeline,
        FlaxStableDiffusionInpaintPipeline,
        FlaxStableDiffusionPipeline,
    )