__init__.py 5.47 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
        UnCLIPImageVariationPipeline,
Will Berman's avatar
Will Berman committed
109
        UnCLIPPipeline,
110
111
112
113
        VersatileDiffusionDualGuidedPipeline,
        VersatileDiffusionImageVariationPipeline,
        VersatileDiffusionPipeline,
        VersatileDiffusionTextToImagePipeline,
114
        VQDiffusionPipeline,
115
    )
116

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

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

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

147
148
149
150
151
152
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from .utils.dummy_flax_objects import *  # noqa F403
else:
153
154
155
156
157
158
159
    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,
160
        FlaxDPMSolverMultistepScheduler,
161
162
163
164
165
166
167
        FlaxKarrasVeScheduler,
        FlaxLMSDiscreteScheduler,
        FlaxPNDMScheduler,
        FlaxSchedulerMixin,
        FlaxScoreSdeVeScheduler,
    )

168

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