__init__.py 4.84 KB
Newer Older
1
# Copyright 2024 The HuggingFace Team. All rights reserved.
2
3
4
5
6
7
8
9
10
11
12
13
14
#
# 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.

15
16
from typing import TYPE_CHECKING

Suraj Patil's avatar
Suraj Patil committed
17
18
19
20
21
22
from ..utils import (
    DIFFUSERS_SLOW_IMPORT,
    _LazyModule,
    is_flax_available,
    is_torch_available,
)
23
24


Dhruv Nair's avatar
Dhruv Nair committed
25
26
_import_structure = {}

27
if is_torch_available():
Dhruv Nair's avatar
Dhruv Nair committed
28
    _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
29
30
31
32
33
    _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"]
    _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"]
    _import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"]
    _import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"]
    _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
34
    _import_structure["autoencoders.vq_model"] = ["VQModel"]
Dhruv Nair's avatar
Dhruv Nair committed
35
    _import_structure["controlnet"] = ["ControlNetModel"]
36
    _import_structure["controlnet_xs"] = ["ControlNetXSAdapter", "UNetControlNetXSModel"]
Dhruv Nair's avatar
Dhruv Nair committed
37
    _import_structure["dual_transformer_2d"] = ["DualTransformer2DModel"]
38
    _import_structure["embeddings"] = ["ImageProjection"]
39
    _import_structure["modeling_utils"] = ["ModelMixin"]
40
    _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"]
41
    _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"]
42
    _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"]
43
44
45
    _import_structure["transformers.prior_transformer"] = ["PriorTransformer"]
    _import_structure["transformers.t5_film_transformer"] = ["T5FilmDecoder"]
    _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
Dhruv Nair's avatar
Dhruv Nair committed
46
    _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"]
47
    _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"]
48
49
50
51
    _import_structure["unets.unet_1d"] = ["UNet1DModel"]
    _import_structure["unets.unet_2d"] = ["UNet2DModel"]
    _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
    _import_structure["unets.unet_3d_condition"] = ["UNet3DConditionModel"]
52
    _import_structure["unets.unet_i2vgen_xl"] = ["I2VGenXLUNet"]
53
54
55
    _import_structure["unets.unet_kandinsky3"] = ["Kandinsky3UNet"]
    _import_structure["unets.unet_motion_model"] = ["MotionAdapter", "UNetMotionModel"]
    _import_structure["unets.unet_spatio_temporal_condition"] = ["UNetSpatioTemporalConditionModel"]
56
    _import_structure["unets.unet_stable_cascade"] = ["StableCascadeUNet"]
57
    _import_structure["unets.uvit_2d"] = ["UVit2DModel"]
58
59

if is_flax_available():
Dhruv Nair's avatar
Dhruv Nair committed
60
    _import_structure["controlnet_flax"] = ["FlaxControlNetModel"]
61
    _import_structure["unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
Dhruv Nair's avatar
Dhruv Nair committed
62
63
64
    _import_structure["vae_flax"] = ["FlaxAutoencoderKL"]


65
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
66
67
    if is_torch_available():
        from .adapter import MultiAdapter, T2IAdapter
68
69
70
71
72
73
        from .autoencoders import (
            AsymmetricAutoencoderKL,
            AutoencoderKL,
            AutoencoderKLTemporalDecoder,
            AutoencoderTiny,
            ConsistencyDecoderVAE,
74
            VQModel,
75
        )
76
        from .controlnet import ControlNetModel
77
        from .controlnet_xs import ControlNetXSAdapter, UNetControlNetXSModel
78
        from .embeddings import ImageProjection
79
        from .modeling_utils import ModelMixin
80
        from .transformers import (
81
            DiTTransformer2DModel,
82
            DualTransformer2DModel,
83
            HunyuanDiT2DModel,
84
            PixArtTransformer2DModel,
85
            PriorTransformer,
Dhruv Nair's avatar
Dhruv Nair committed
86
            SD3Transformer2DModel,
87
88
89
90
            T5FilmDecoder,
            Transformer2DModel,
            TransformerTemporalModel,
        )
91
        from .unets import (
92
            I2VGenXLUNet,
93
94
            Kandinsky3UNet,
            MotionAdapter,
95
            StableCascadeUNet,
96
97
98
99
100
101
102
103
            UNet1DModel,
            UNet2DConditionModel,
            UNet2DModel,
            UNet3DConditionModel,
            UNetMotionModel,
            UNetSpatioTemporalConditionModel,
            UVit2DModel,
        )
104
105
106

    if is_flax_available():
        from .controlnet_flax import FlaxControlNetModel
107
        from .unets import FlaxUNet2DConditionModel
108
109
110
111
        from .vae_flax import FlaxAutoencoderKL

else:
    import sys
Dhruv Nair's avatar
Dhruv Nair committed
112

113
    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)