__init__.py 5.4 KB
Newer Older
Anton Lozhkov's avatar
Anton Lozhkov committed
1
__version__ = "0.12.0.dev0"
2
3
4

from .configuration_utils import ConfigMixin
from .onnx_utils import OnnxRuntimeModel
5
from .utils import (
6
    OptionalDependencyNotAvailable,
7
    is_flax_available,
8
    is_inflect_available,
9
    is_k_diffusion_available,
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
try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_pt_objects import *  # noqa F403
else:
27
    from .modeling_utils import ModelMixin
Will Berman's avatar
Will Berman committed
28
29
30
31
32
33
34
35
36
    from .models import (
        AutoencoderKL,
        PriorTransformer,
        Transformer2DModel,
        UNet1DModel,
        UNet2DConditionModel,
        UNet2DModel,
        VQModel,
    )
37
38
39
40
41
42
43
44
45
46
    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,
    )
    from .pipeline_utils import DiffusionPipeline
47
48
49
50
51
52
    from .pipelines import (
        DanceDiffusionPipeline,
        DDIMPipeline,
        DDPMPipeline,
        KarrasVePipeline,
        LDMPipeline,
53
        LDMSuperResolutionPipeline,
54
        PNDMPipeline,
Revist's avatar
Revist committed
55
        RePaintPipeline,
56
57
        ScoreSdeVePipeline,
    )
58
59
60
    from .schedulers import (
        DDIMScheduler,
        DDPMScheduler,
61
        DPMSolverMultistepScheduler,
62
        DPMSolverSinglestepScheduler,
hlky's avatar
hlky committed
63
64
        EulerAncestralDiscreteScheduler,
        EulerDiscreteScheduler,
65
        HeunDiscreteScheduler,
66
        IPNDMScheduler,
67
        KarrasVeScheduler,
68
69
        KDPM2AncestralDiscreteScheduler,
        KDPM2DiscreteScheduler,
70
        PNDMScheduler,
Revist's avatar
Revist committed
71
        RePaintScheduler,
72
73
        SchedulerMixin,
        ScoreSdeVeScheduler,
Will Berman's avatar
Will Berman committed
74
        UnCLIPScheduler,
Will Berman's avatar
Will Berman committed
75
        VQDiffusionScheduler,
76
77
    )
    from .training_utils import EMAModel
Patrick von Platen's avatar
Patrick von Platen committed
78

79
80
81
82
try:
    if not (is_torch_available() and is_scipy_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
83
    from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
84
85
else:
    from .schedulers import LMSDiscreteScheduler
Patrick von Platen's avatar
Patrick von Platen committed
86

87
88
89
90
91
92
93

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:
94
    from .pipelines import (
Patrick von Platen's avatar
Patrick von Platen committed
95
96
        AltDiffusionImg2ImgPipeline,
        AltDiffusionPipeline,
97
        CycleDiffusionPipeline,
98
        LDMTextToImagePipeline,
99
        PaintByExamplePipeline,
100
        StableDiffusionDepth2ImgPipeline,
101
        StableDiffusionImageVariationPipeline,
102
103
        StableDiffusionImg2ImgPipeline,
        StableDiffusionInpaintPipeline,
104
        StableDiffusionInpaintPipelineLegacy,
105
        StableDiffusionPipeline,
106
        StableDiffusionPipelineSafe,
107
        StableDiffusionUpscalePipeline,
Will Berman's avatar
Will Berman committed
108
        UnCLIPPipeline,
109
110
111
112
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
113
        VQDiffusionPipeline,
114
    )
115

116
117
118
119
try:
    if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
120
    from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
121
122
else:
    from .pipelines import StableDiffusionKDiffusionPipeline
123

124
125
126
127
128
129
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:
130
131
132
    from .pipelines import (
        OnnxStableDiffusionImg2ImgPipeline,
        OnnxStableDiffusionInpaintPipeline,
133
        OnnxStableDiffusionInpaintPipelineLegacy,
134
135
136
        OnnxStableDiffusionPipeline,
        StableDiffusionOnnxPipeline,
    )
137
138
139
140
141
142

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

146
147
148
149
150
151
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
152
153
154
155
156
157
158
    from .modeling_flax_utils import FlaxModelMixin
    from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
    from .models.vae_flax import FlaxAutoencoderKL
    from .pipeline_flax_utils import FlaxDiffusionPipeline
    from .schedulers import (
        FlaxDDIMScheduler,
        FlaxDDPMScheduler,
159
        FlaxDPMSolverMultistepScheduler,
160
161
162
163
164
165
166
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

167
168
169
170
try:
    if not (is_flax_available() and is_transformers_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
171
    from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
172
173
else:
    from .pipelines import FlaxStableDiffusionPipeline