__init__.py 10.9 KB
Newer Older
1
# Copyright 2024 The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
14

15
from typing import TYPE_CHECKING
16

17
from ..utils import (
18
    DIFFUSERS_SLOW_IMPORT,
19
    OptionalDependencyNotAvailable,
Dhruv Nair's avatar
Dhruv Nair committed
20
    _LazyModule,
21
    get_objects_from_module,
22
23
24
25
26
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)
27
28


Dhruv Nair's avatar
Dhruv Nair committed
29
_dummy_modules = {}
30
_import_structure = {}
Dhruv Nair's avatar
Dhruv Nair committed
31

32
33
34
35
try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
Dhruv Nair's avatar
Dhruv Nair committed
36
37
    from ..utils import dummy_pt_objects  # noqa F403

38
    _dummy_modules.update(get_objects_from_module(dummy_pt_objects))
Dhruv Nair's avatar
Dhruv Nair committed
39

40
else:
41
    _import_structure["deprecated"] = ["KarrasVeScheduler", "ScoreSdeVpScheduler"]
Will Berman's avatar
Will Berman committed
42
    _import_structure["scheduling_amused"] = ["AmusedScheduler"]
Will Berman's avatar
Will Berman committed
43
    _import_structure["scheduling_consistency_decoder"] = ["ConsistencyDecoderScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
44
45
    _import_structure["scheduling_consistency_models"] = ["CMStochasticIterativeScheduler"]
    _import_structure["scheduling_ddim"] = ["DDIMScheduler"]
46
    _import_structure["scheduling_ddim_cogvideox"] = ["CogVideoXDDIMScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
47
48
49
50
    _import_structure["scheduling_ddim_inverse"] = ["DDIMInverseScheduler"]
    _import_structure["scheduling_ddim_parallel"] = ["DDIMParallelScheduler"]
    _import_structure["scheduling_ddpm"] = ["DDPMScheduler"]
    _import_structure["scheduling_ddpm_parallel"] = ["DDPMParallelScheduler"]
51
    _import_structure["scheduling_ddpm_wuerstchen"] = ["DDPMWuerstchenScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
52
    _import_structure["scheduling_deis_multistep"] = ["DEISMultistepScheduler"]
53
    _import_structure["scheduling_dpm_cogvideox"] = ["CogVideoXDPMScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
54
55
56
    _import_structure["scheduling_dpmsolver_multistep"] = ["DPMSolverMultistepScheduler"]
    _import_structure["scheduling_dpmsolver_multistep_inverse"] = ["DPMSolverMultistepInverseScheduler"]
    _import_structure["scheduling_dpmsolver_singlestep"] = ["DPMSolverSinglestepScheduler"]
57
    _import_structure["scheduling_edm_dpmsolver_multistep"] = ["EDMDPMSolverMultistepScheduler"]
Suraj Patil's avatar
Suraj Patil committed
58
    _import_structure["scheduling_edm_euler"] = ["EDMEulerScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
59
60
    _import_structure["scheduling_euler_ancestral_discrete"] = ["EulerAncestralDiscreteScheduler"]
    _import_structure["scheduling_euler_discrete"] = ["EulerDiscreteScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
61
    _import_structure["scheduling_flow_match_euler_discrete"] = ["FlowMatchEulerDiscreteScheduler"]
62
    _import_structure["scheduling_flow_match_heun_discrete"] = ["FlowMatchHeunDiscreteScheduler"]
63
    _import_structure["scheduling_flow_match_lcm"] = ["FlowMatchLCMScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
64
65
66
67
    _import_structure["scheduling_heun_discrete"] = ["HeunDiscreteScheduler"]
    _import_structure["scheduling_ipndm"] = ["IPNDMScheduler"]
    _import_structure["scheduling_k_dpm_2_ancestral_discrete"] = ["KDPM2AncestralDiscreteScheduler"]
    _import_structure["scheduling_k_dpm_2_discrete"] = ["KDPM2DiscreteScheduler"]
68
    _import_structure["scheduling_lcm"] = ["LCMScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
69
70
    _import_structure["scheduling_pndm"] = ["PNDMScheduler"]
    _import_structure["scheduling_repaint"] = ["RePaintScheduler"]
Junsong Chen's avatar
Junsong Chen committed
71
    _import_structure["scheduling_sasolver"] = ["SASolverScheduler"]
YiYi Xu's avatar
YiYi Xu committed
72
    _import_structure["scheduling_scm"] = ["SCMScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
73
    _import_structure["scheduling_sde_ve"] = ["ScoreSdeVeScheduler"]
Michael's avatar
Michael committed
74
    _import_structure["scheduling_tcd"] = ["TCDScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
75
76
    _import_structure["scheduling_unclip"] = ["UnCLIPScheduler"]
    _import_structure["scheduling_unipc_multistep"] = ["UniPCMultistepScheduler"]
77
    _import_structure["scheduling_utils"] = ["AysSchedules", "KarrasDiffusionSchedulers", "SchedulerMixin"]
Dhruv Nair's avatar
Dhruv Nair committed
78
    _import_structure["scheduling_vq_diffusion"] = ["VQDiffusionScheduler"]
79

80
81
82
83
try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
84
85
86
87
    from ..utils import dummy_flax_objects  # noqa F403

    _dummy_modules.update(get_objects_from_module(dummy_flax_objects))

88
else:
Dhruv Nair's avatar
Dhruv Nair committed
89
90
91
    _import_structure["scheduling_ddim_flax"] = ["FlaxDDIMScheduler"]
    _import_structure["scheduling_ddpm_flax"] = ["FlaxDDPMScheduler"]
    _import_structure["scheduling_dpmsolver_multistep_flax"] = ["FlaxDPMSolverMultistepScheduler"]
Pedro Cuenca's avatar
Pedro Cuenca committed
92
    _import_structure["scheduling_euler_discrete_flax"] = ["FlaxEulerDiscreteScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
93
94
95
96
97
98
99
100
101
102
    _import_structure["scheduling_karras_ve_flax"] = ["FlaxKarrasVeScheduler"]
    _import_structure["scheduling_lms_discrete_flax"] = ["FlaxLMSDiscreteScheduler"]
    _import_structure["scheduling_pndm_flax"] = ["FlaxPNDMScheduler"]
    _import_structure["scheduling_sde_ve_flax"] = ["FlaxScoreSdeVeScheduler"]
    _import_structure["scheduling_utils_flax"] = [
        "FlaxKarrasDiffusionSchedulers",
        "FlaxSchedulerMixin",
        "FlaxSchedulerOutput",
        "broadcast_to_shape_from_left",
    ]
103

104

105
106
107
108
try:
    if not (is_torch_available() and is_scipy_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
Dhruv Nair's avatar
Dhruv Nair committed
109
110
    from ..utils import dummy_torch_and_scipy_objects  # noqa F403

111
    _dummy_modules.update(get_objects_from_module(dummy_torch_and_scipy_objects))
Dhruv Nair's avatar
Dhruv Nair committed
112

113
else:
Dhruv Nair's avatar
Dhruv Nair committed
114
    _import_structure["scheduling_lms_discrete"] = ["LMSDiscreteScheduler"]
115
116
117
118
119

try:
    if not (is_torch_available() and is_torchsde_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
Dhruv Nair's avatar
Dhruv Nair committed
120
121
    from ..utils import dummy_torch_and_torchsde_objects  # noqa F403

122
    _dummy_modules.update(get_objects_from_module(dummy_torch_and_torchsde_objects))
Dhruv Nair's avatar
Dhruv Nair committed
123

124
else:
125
    _import_structure["scheduling_cosine_dpmsolver_multistep"] = ["CosineDPMSolverMultistepScheduler"]
Dhruv Nair's avatar
Dhruv Nair committed
126
127
    _import_structure["scheduling_dpmsolver_sde"] = ["DPMSolverSDEScheduler"]

128
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
129
130
131
132
133
134
135
136
137
138
139
140
141
142
    from ..utils import (
        OptionalDependencyNotAvailable,
        is_flax_available,
        is_scipy_available,
        is_torch_available,
        is_torchsde_available,
    )

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        from ..utils.dummy_pt_objects import *  # noqa F403
    else:
143
        from .deprecated import KarrasVeScheduler, ScoreSdeVpScheduler
Will Berman's avatar
Will Berman committed
144
        from .scheduling_amused import AmusedScheduler
Will Berman's avatar
Will Berman committed
145
        from .scheduling_consistency_decoder import ConsistencyDecoderScheduler
146
147
        from .scheduling_consistency_models import CMStochasticIterativeScheduler
        from .scheduling_ddim import DDIMScheduler
148
        from .scheduling_ddim_cogvideox import CogVideoXDDIMScheduler
149
150
151
152
153
154
        from .scheduling_ddim_inverse import DDIMInverseScheduler
        from .scheduling_ddim_parallel import DDIMParallelScheduler
        from .scheduling_ddpm import DDPMScheduler
        from .scheduling_ddpm_parallel import DDPMParallelScheduler
        from .scheduling_ddpm_wuerstchen import DDPMWuerstchenScheduler
        from .scheduling_deis_multistep import DEISMultistepScheduler
155
        from .scheduling_dpm_cogvideox import CogVideoXDPMScheduler
156
157
158
        from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
        from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
        from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
159
        from .scheduling_edm_dpmsolver_multistep import EDMDPMSolverMultistepScheduler
Suraj Patil's avatar
Suraj Patil committed
160
        from .scheduling_edm_euler import EDMEulerScheduler
161
162
        from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
        from .scheduling_euler_discrete import EulerDiscreteScheduler
Dhruv Nair's avatar
Dhruv Nair committed
163
        from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
164
        from .scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler
165
        from .scheduling_flow_match_lcm import FlowMatchLCMScheduler
166
167
168
169
        from .scheduling_heun_discrete import HeunDiscreteScheduler
        from .scheduling_ipndm import IPNDMScheduler
        from .scheduling_k_dpm_2_ancestral_discrete import KDPM2AncestralDiscreteScheduler
        from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler
170
        from .scheduling_lcm import LCMScheduler
171
172
        from .scheduling_pndm import PNDMScheduler
        from .scheduling_repaint import RePaintScheduler
Junsong Chen's avatar
Junsong Chen committed
173
        from .scheduling_sasolver import SASolverScheduler
YiYi Xu's avatar
YiYi Xu committed
174
        from .scheduling_scm import SCMScheduler
175
        from .scheduling_sde_ve import ScoreSdeVeScheduler
Michael's avatar
Michael committed
176
        from .scheduling_tcd import TCDScheduler
177
178
        from .scheduling_unclip import UnCLIPScheduler
        from .scheduling_unipc_multistep import UniPCMultistepScheduler
179
        from .scheduling_utils import AysSchedules, KarrasDiffusionSchedulers, SchedulerMixin
180
181
182
183
184
185
186
187
188
189
        from .scheduling_vq_diffusion import VQDiffusionScheduler
    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        from ..utils.dummy_flax_objects import *  # noqa F403
    else:
        from .scheduling_ddim_flax import FlaxDDIMScheduler
        from .scheduling_ddpm_flax import FlaxDDPMScheduler
        from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
190
        from .scheduling_euler_discrete_flax import FlaxEulerDiscreteScheduler
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
        from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
        from .scheduling_pndm_flax import FlaxPNDMScheduler
        from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
        from .scheduling_utils_flax import (
            FlaxKarrasDiffusionSchedulers,
            FlaxSchedulerMixin,
            FlaxSchedulerOutput,
            broadcast_to_shape_from_left,
        )

    try:
        if not (is_torch_available() and is_scipy_available()):
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
    else:
        from .scheduling_lms_discrete import LMSDiscreteScheduler

    try:
        if not (is_torch_available() and is_torchsde_available()):
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
    else:
216
        from .scheduling_cosine_dpmsolver_multistep import CosineDPMSolverMultistepScheduler
217
        from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler
Dhruv Nair's avatar
Dhruv Nair committed
218

219
220
else:
    import sys
Dhruv Nair's avatar
Dhruv Nair committed
221

222
223
224
    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
    for name, value in _dummy_modules.items():
        setattr(sys.modules[__name__], name, value)