__init__.py 2.66 KB
Newer Older
YiYi Xu's avatar
YiYi Xu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
from typing import TYPE_CHECKING

from ..utils import (
    DIFFUSERS_SLOW_IMPORT,
    OptionalDependencyNotAvailable,
    _LazyModule,
    get_objects_from_module,
    is_torch_available,
    is_transformers_available,
)


# These modules contain pipelines from multiple libraries/frameworks
_dummy_objects = {}
_import_structure = {}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ..utils import dummy_pt_objects  # noqa F403

    _dummy_objects.update(get_objects_from_module(dummy_pt_objects))
else:
    _import_structure["modular_pipeline"] = [
        "ModularPipelineBlocks",
        "ModularPipeline",
        "PipelineBlock",
        "AutoPipelineBlocks",
        "SequentialPipelineBlocks",
        "LoopSequentialPipelineBlocks",
        "PipelineState",
        "BlockState",
    ]
    _import_structure["modular_pipeline_utils"] = [
        "ComponentSpec",
        "ConfigSpec",
        "InputParam",
        "OutputParam",
        "InsertableDict",
    ]
    _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"]
Aryan's avatar
Aryan committed
43
    _import_structure["wan"] = ["WanAutoBlocks", "WanModularPipeline"]
44
    _import_structure["flux"] = ["FluxAutoBlocks", "FluxModularPipeline"]
YiYi Xu's avatar
YiYi Xu committed
45
46
47
48
49
50
51
52
53
54
    _import_structure["components_manager"] = ["ComponentsManager"]

if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        from ..utils.dummy_pt_objects import *  # noqa F403
    else:
        from .components_manager import ComponentsManager
55
        from .flux import FluxAutoBlocks, FluxModularPipeline
YiYi Xu's avatar
YiYi Xu committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
        from .modular_pipeline import (
            AutoPipelineBlocks,
            BlockState,
            LoopSequentialPipelineBlocks,
            ModularPipeline,
            ModularPipelineBlocks,
            PipelineBlock,
            PipelineState,
            SequentialPipelineBlocks,
        )
        from .modular_pipeline_utils import (
            ComponentSpec,
            ConfigSpec,
            InputParam,
            InsertableDict,
            OutputParam,
        )
        from .stable_diffusion_xl import (
            StableDiffusionXLAutoBlocks,
            StableDiffusionXLModularPipeline,
        )
Aryan's avatar
Aryan committed
77
        from .wan import WanAutoBlocks, WanModularPipeline
YiYi Xu's avatar
YiYi Xu committed
78
79
80
81
82
83
84
85
86
87
88
else:
    import sys

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