__init__.py 6.05 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
        StableDiffusionPix2PixZeroPipeline,
122
        StableDiffusionSAGPipeline,
123
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
124
125
        StableUnCLIPImg2ImgPipeline,
        StableUnCLIPPipeline,
Will Berman's avatar
Will Berman committed
126
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
127
        UnCLIPPipeline,
128
129
130
131
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
132
        VQDiffusionPipeline,
133
    )
134

135
136
137
138
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
139
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
140
141
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
142

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

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

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

186

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