diffusers/modular_pipelines/__init__.py

from typing import TYPE_CHECKING

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


logger = logging.get_logger(__name__)
logger.warning(
    "Modular Diffusers is currently an experimental feature under active development. The API is subject to breaking changes in future releases."
)

# 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",
        "AutoPipelineBlocks",
        "SequentialPipelineBlocks",
        "ConditionalPipelineBlocks",
        "LoopSequentialPipelineBlocks",
        "PipelineState",
        "BlockState",
    ]
    _import_structure["modular_pipeline_utils"] = [
        "ComponentSpec",
        "ConfigSpec",
        "InputParam",
        "OutputParam",
        "InsertableDict",
    ]
    _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"]
    _import_structure["wan"] = [
        "WanBlocks",
        "Wan22Blocks",
        "WanImage2VideoAutoBlocks",
        "Wan22Image2VideoBlocks",
        "WanModularPipeline",
        "Wan22ModularPipeline",
        "WanImage2VideoModularPipeline",
        "Wan22Image2VideoModularPipeline",
    ]
    _import_structure["flux"] = [
        "FluxAutoBlocks",
        "FluxModularPipeline",
        "FluxKontextAutoBlocks",
        "FluxKontextModularPipeline",
    ]
    _import_structure["flux2"] = [
        "Flux2AutoBlocks",
        "Flux2KleinAutoBlocks",
        "Flux2KleinBaseAutoBlocks",
        "Flux2ModularPipeline",
        "Flux2KleinModularPipeline",
        "Flux2KleinBaseModularPipeline",
    ]
    _import_structure["qwenimage"] = [
        "QwenImageAutoBlocks",
        "QwenImageModularPipeline",
        "QwenImageEditModularPipeline",
        "QwenImageEditAutoBlocks",
        "QwenImageEditPlusModularPipeline",
        "QwenImageEditPlusAutoBlocks",
        "QwenImageLayeredModularPipeline",
        "QwenImageLayeredAutoBlocks",
    ]
    _import_structure["z_image"] = [
        "ZImageAutoBlocks",
        "ZImageModularPipeline",
    ]
    _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
        from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
        from .flux2 import (
            Flux2AutoBlocks,
            Flux2KleinAutoBlocks,
            Flux2KleinBaseAutoBlocks,
            Flux2KleinBaseModularPipeline,
            Flux2KleinModularPipeline,
            Flux2ModularPipeline,
        )
        from .modular_pipeline import (
            AutoPipelineBlocks,
            BlockState,
            ConditionalPipelineBlocks,
            LoopSequentialPipelineBlocks,
            ModularPipeline,
            ModularPipelineBlocks,
            PipelineState,
            SequentialPipelineBlocks,
        )
        from .modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, InsertableDict, OutputParam
        from .qwenimage import (
            QwenImageAutoBlocks,
            QwenImageEditAutoBlocks,
            QwenImageEditModularPipeline,
            QwenImageEditPlusAutoBlocks,
            QwenImageEditPlusModularPipeline,
            QwenImageLayeredAutoBlocks,
            QwenImageLayeredModularPipeline,
            QwenImageModularPipeline,
        )
        from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline
        from .wan import (
            Wan22Blocks,
            Wan22Image2VideoBlocks,
            Wan22Image2VideoModularPipeline,
            Wan22ModularPipeline,
            WanBlocks,
            WanImage2VideoAutoBlocks,
            WanImage2VideoModularPipeline,
            WanModularPipeline,
        )
        from .z_image import ZImageAutoBlocks, ZImageModularPipeline
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)