cache-dit 1.0.3__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_pixart.py

@@ -0,0 +1,285 @@
+import torch
+import functools
+from typing import Optional
+import torch.nn.functional as F
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformers.pixart_transformer_2d import (
+    PixArtTransformer2DModel,
+)
+from diffusers.models.attention_processor import (
+    Attention,
+    AttnProcessor2_0,
+)  # sdpa
+from diffusers.utils import deprecate
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("PixArt")
+class PixArtContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        assert transformer is not None, "Transformer must be provided."
+        assert isinstance(
+            transformer, PixArtTransformer2DModel
+        ), "Transformer must be an instance of PixArtTransformer2DModel"
+
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Apply monkey patch to fix attention mask preparation at class level
+        Attention.prepare_attention_mask = (
+            __patch_Attention_prepare_attention_mask__
+        )
+        AttnProcessor2_0.__call__ = __patch_AttnProcessor2_0__call__
+        if not hasattr(AttnProcessor2_0, "_parallel_config"):
+            AttnProcessor2_0._parallel_config = None
+        if not hasattr(AttnProcessor2_0, "_attention_backend"):
+            AttnProcessor2_0._attention_backend = None
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+
+        _cp_plan = {
+            # Pattern of transformer_blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            #     (only split hidden_states, not encoder_hidden_states)
+            "transformer_blocks.0": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            # Pattern of the all blocks, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            #    (only split encoder_hidden_states, not hidden_states.
+            #    hidden_states has been automatically split in previous
+            #    block by all2all comm op after attn)
+            # The `encoder_hidden_states` will [NOT] be changed after each block forward,
+            # so we need to split it at [ALL] block by the inserted split hook.
+            "transformer_blocks.*": {
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan
+
+
+@functools.wraps(Attention.prepare_attention_mask)
+def __patch_Attention_prepare_attention_mask__(
+    self: Attention,
+    attention_mask: torch.Tensor,
+    target_length: int,
+    batch_size: int,
+    out_dim: int = 3,
+    # NOTE(DefTruth): Allow specifying head_size for CP
+    head_size: Optional[int] = None,
+) -> torch.Tensor:
+    if head_size is None:
+        head_size = self.heads
+    if attention_mask is None:
+        return attention_mask
+
+    current_length: int = attention_mask.shape[-1]
+    if current_length != target_length:
+        if attention_mask.device.type == "mps":
+            # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+            # Instead, we can manually construct the padding tensor.
+            padding_shape = (
+                attention_mask.shape[0],
+                attention_mask.shape[1],
+                target_length,
+            )
+            padding = torch.zeros(
+                padding_shape,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+            attention_mask = torch.cat([attention_mask, padding], dim=2)
+        else:
+            # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+            #       we want to instead pad by (0, remaining_length), where remaining_length is:
+            #       remaining_length: int = target_length - current_length
+            # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+            attention_mask = F.pad(
+                attention_mask, (0, target_length), value=0.0
+            )
+
+    if out_dim == 3:
+        if attention_mask.shape[0] < batch_size * head_size:
+            attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+    elif out_dim == 4:
+        attention_mask = attention_mask.unsqueeze(1)
+        attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+    return attention_mask
+
+
+@functools.wraps(AttnProcessor2_0.__call__)
+def __patch_AttnProcessor2_0__call__(
+    self: AttnProcessor2_0,
+    attn: Attention,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    temb: Optional[torch.Tensor] = None,
+    *args,
+    **kwargs,
+) -> torch.Tensor:
+    if len(args) > 0 or kwargs.get("scale", None) is not None:
+        deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+        deprecate("scale", "1.0.0", deprecation_message)
+
+    residual = hidden_states
+    if attn.spatial_norm is not None:
+        hidden_states = attn.spatial_norm(hidden_states, temb)
+
+    input_ndim = hidden_states.ndim
+
+    if input_ndim == 4:
+        batch_size, channel, height, width = hidden_states.shape
+        hidden_states = hidden_states.view(
+            batch_size, channel, height * width
+        ).transpose(1, 2)
+
+    batch_size, sequence_length, _ = (
+        hidden_states.shape
+        if encoder_hidden_states is None
+        else encoder_hidden_states.shape
+    )
+
+    if attention_mask is not None:
+        if self._parallel_config is None:
+            attention_mask = attn.prepare_attention_mask(
+                attention_mask, sequence_length, batch_size
+            )
+            attention_mask = attention_mask.view(
+                batch_size, attn.heads, -1, attention_mask.shape[-1]
+            )
+        else:
+            # NOTE(DefTruth): Fix attention mask preparation for context parallelism.
+            # Please note that in context parallelism, the sequence_length is the local
+            # sequence length on each rank. So we need to adjust the target_length
+            # accordingly. The head_size is also adjusted based on the world size
+            # in order to make sdpa work correctly, otherwise, the sdpa op will raise
+            # error due to the mismatch between attention_mask shape and expected shape.
+            cp_config = getattr(
+                self._parallel_config, "context_parallel_config", None
+            )
+            if cp_config is not None and cp_config._world_size > 1:
+                head_size = attn.heads // cp_config._world_size
+                attention_mask = attn.prepare_attention_mask(
+                    attention_mask,
+                    sequence_length * cp_config._world_size,
+                    batch_size,
+                    3,
+                    head_size,
+                )
+                attention_mask = attention_mask.view(
+                    batch_size, head_size, -1, attention_mask.shape[-1]
+                )
+
+    if attn.group_norm is not None:
+        hidden_states = attn.group_norm(
+            hidden_states.transpose(1, 2)
+        ).transpose(1, 2)
+
+    query = attn.to_q(hidden_states)
+
+    if encoder_hidden_states is None:
+        encoder_hidden_states = hidden_states
+    elif attn.norm_cross:
+        encoder_hidden_states = attn.norm_encoder_hidden_states(
+            encoder_hidden_states
+        )
+
+    key = attn.to_k(encoder_hidden_states)
+    value = attn.to_v(encoder_hidden_states)
+
+    inner_dim = key.shape[-1]
+    head_dim = inner_dim // attn.heads
+
+    # NOTE(DefTruth): no transpose now
+    query = query.view(batch_size, -1, attn.heads, head_dim)
+    key = key.view(batch_size, -1, attn.heads, head_dim)
+    value = value.view(batch_size, -1, attn.heads, head_dim)
+
+    if attn.norm_q is not None:
+        query = attn.norm_q(query)
+    if attn.norm_k is not None:
+        key = attn.norm_k(key)
+
+    # NOTE(DefTruth): Use the dispatch_attention_fn to support different backends
+    hidden_states = dispatch_attention_fn(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=0.0,
+        is_causal=False,
+        backend=getattr(self, "_attention_backend", None),
+        parallel_config=getattr(self, "_parallel_config", None),
+    )
+    hidden_states = hidden_states.flatten(2, 3)
+    hidden_states = hidden_states.to(query.dtype)
+
+    # linear proj
+    hidden_states = attn.to_out[0](hidden_states)
+    # dropout
+    hidden_states = attn.to_out[1](hidden_states)
+
+    if input_ndim == 4:
+        hidden_states = hidden_states.transpose(-1, -2).reshape(
+            batch_size, channel, height, width
+        )
+
+    if attn.residual_connection:
+        hidden_states = hidden_states + residual
+
+    hidden_states = hidden_states / attn.rescale_output_factor
+
+    return hidden_states

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_qwen_image.py

@@ -0,0 +1,104 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@ContextParallelismPlannerRegister.register("QwenImage")
+class QwenImageContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+
+        # NOTE: Set it as False to use custom CP plan defined here.
+        self._cp_planner_preferred_native_diffusers = False
+
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            from diffusers import QwenImageTransformer2DModel
+
+            assert isinstance(
+                transformer, QwenImageTransformer2DModel
+            ), "Transformer must be an instance of QwenImageTransformer2DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Here is a Transformer level CP plan for Flux, which will
+            # only apply the only 1 split hook (pre_forward) on the forward
+            # of Transformer, and gather the output after Transformer forward.
+            # Pattern of transformer forward, split_output=False:
+            #     un-split input -> splited input (inside transformer)
+            # Pattern of the transformer_blocks, single_transformer_blocks:
+            #     splited input (previous splited output) -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            # The `hidden_states` and `encoder_hidden_states` will still keep
+            # itself splited after block forward (namely, automatic split by
+            # the all2all comm op after attn) for the all blocks.
+            "": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                # NOTE: Due to the joint attention implementation of
+                # QwenImageTransformerBlock, we must split the
+                # encoder_hidden_states as well.
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+                # NOTE: But encoder_hidden_states_mask seems never used in
+                # QwenImageTransformerBlock, so we do not split it here.
+                # "encoder_hidden_states_mask": ContextParallelInput(
+                #     split_dim=1, expected_dims=2, split_output=False
+                # ),
+            },
+            # Pattern of pos_embed, split_output=True (split output rather than input):
+            #    un-split input
+            #    -> keep input un-split
+            #    -> rope
+            #    -> splited output
+            "pos_embed": {
+                0: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=0, expected_dims=2, split_output=True
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_registers.py

@@ -0,0 +1,84 @@
+import torch
+import logging
+from abc import abstractmethod
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+__all__ = [
+    "ContextParallelismPlanner",
+    "ContextParallelismPlannerRegister",
+]
+
+
+class ContextParallelismPlanner:
+    # Prefer native diffusers implementation if available
+    _cp_planner_preferred_native_diffusers: bool = True
+
+    @abstractmethod
+    def apply(
+        self,
+        # NOTE: Keep this kwarg for future extensions
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        # NOTE: This method should only return the CP plan dictionary.
+        raise NotImplementedError(
+            "apply method must be implemented by subclasses"
+        )
+
+
+class ContextParallelismPlannerRegister:
+    _cp_planner_registry: dict[str, ContextParallelismPlanner] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        def decorator(planner_cls: type[ContextParallelismPlanner]):
+            assert (
+                name not in cls._cp_planner_registry
+            ), f"ContextParallelismPlanner with name {name} is already registered."
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(f"Registering ContextParallelismPlanner: {name}")
+            cls._cp_planner_registry[name] = planner_cls
+            return planner_cls
+
+        return decorator
+
+    @classmethod
+    def get_planner(
+        cls, transformer: str | torch.nn.Module | ModelMixin
+    ) -> type[ContextParallelismPlanner]:
+        if isinstance(transformer, (torch.nn.Module, ModelMixin)):
+            name = transformer.__class__.__name__
+        else:
+            name = transformer
+        planner_cls = None
+        for planner_name in cls._cp_planner_registry:
+            if name.startswith(planner_name):
+                planner_cls = cls._cp_planner_registry.get(planner_name)
+                break
+        if planner_cls is None:
+            raise ValueError(f"No planner registered under name: {name}")
+        return planner_cls
+
+    @classmethod
+    def supported_planners(
+        cls,
+    ) -> tuple[int, list[str]]:
+        val_planners = cls._cp_planner_registry.keys()
+        return len(val_planners), [p for p in val_planners]

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_plan_wan.py

@@ -0,0 +1,101 @@
+import torch
+from typing import Optional
+from diffusers.models.modeling_utils import ModelMixin
+
+try:
+    from diffusers.models._modeling_parallel import (
+        ContextParallelInput,
+        ContextParallelOutput,
+        ContextParallelModelPlan,
+    )
+except ImportError:
+    raise ImportError(
+        "Context parallelism requires the 'diffusers>=0.36.dev0'."
+        "Please install latest version of diffusers from source: \n"
+        "pip3 install git+https://github.com/huggingface/diffusers.git"
+    )
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# TODO: Add WanVACETransformer3DModel context parallelism planner.
+# NOTE: Maybe use full name to avoid name conflict between
+# WanTransformer3DModel and WanVACETransformer3DModel?
+@ContextParallelismPlannerRegister.register("WanTransformer3D")
+class WanContextParallelismPlanner(ContextParallelismPlanner):
+    def apply(
+        self,
+        transformer: Optional[torch.nn.Module | ModelMixin] = None,
+        **kwargs,
+    ) -> ContextParallelModelPlan:
+        if (
+            transformer is not None
+            and self._cp_planner_preferred_native_diffusers
+        ):
+            from diffusers import WanTransformer3DModel
+
+            assert isinstance(
+                transformer, WanTransformer3DModel
+            ), "Transformer must be an instance of WanTransformer3DModel"
+            if hasattr(transformer, "_cp_plan"):
+                if transformer._cp_plan is not None:
+                    return transformer._cp_plan
+
+        # Otherwise, use the custom CP plan defined here, this maybe
+        # a little different from the native diffusers implementation
+        # for some models.
+        _cp_plan = {
+            # Pattern of rope, split_output=True (split output rather than input):
+            #    un-split input
+            #    -> keep input un-split
+            #    -> rope
+            #    -> splited output
+            "rope": {
+                0: ContextParallelInput(
+                    split_dim=1, expected_dims=4, split_output=True
+                ),
+                1: ContextParallelInput(
+                    split_dim=1, expected_dims=4, split_output=True
+                ),
+            },
+            # Pattern of blocks.0, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            #     (only split hidden_states, not encoder_hidden_states)
+            "blocks.0": {
+                "hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            # Pattern of the all blocks, split_output=False:
+            #     un-split input -> split -> to_qkv/...
+            #     -> all2all
+            #     -> attn (local head, full seqlen)
+            #     -> all2all
+            #     -> splited output
+            #    (only split encoder_hidden_states, not hidden_states.
+            #    hidden_states has been automatically split in previous
+            #    block by all2all comm op after attn)
+            # The `encoder_hidden_states` will [NOT] be changed after each block forward,
+            # so we need to split it at [ALL] block by the inserted split hook.
+            "blocks.*": {
+                "encoder_hidden_states": ContextParallelInput(
+                    split_dim=1, expected_dims=3, split_output=False
+                ),
+            },
+            # Then, the final proj_out will gather the splited output.
+            #     splited input (previous splited output)
+            #     -> all gather
+            #     -> un-split output
+            "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+        }
+        return _cp_plan

cache_dit/parallelism/backends/native_diffusers/context_parallelism/cp_planners.py

@@ -0,0 +1,117 @@
+# Docstring references: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py#L185
+# A dictionary where keys denote the input to be split across context parallel region, and the
+# value denotes the sharding configuration.
+# If the key is a string, it denotes the name of the parameter in the forward function.
+# If the key is an integer, split_output must be set to True, and it denotes the index of the output
+# to be split across context parallel region.
+# ContextParallelInputType = Dict[
+#     Union[str, int], Union[ContextParallelInput, List[ContextParallelInput], Tuple[ContextParallelInput, ...]]
+# ]
+
+# A dictionary where keys denote the output to be gathered across context parallel region, and the
+# value denotes the gathering configuration.
+# ContextParallelOutputType = Union[
+#     ContextParallelOutput, List[ContextParallelOutput], Tuple[ContextParallelOutput, ...]
+# ]
+
+# A dictionary where keys denote the module id, and the value denotes how the inputs/outputs of
+# the module should be split/gathered across context parallel region.
+# ContextParallelModelPlan = Dict[str, Union[ContextParallelInputType, ContextParallelOutputType]]
+
+# Example of a ContextParallelModelPlan (QwenImageTransformer2DModel):
+#
+# Each model should define a _cp_plan attribute that contains information on how to shard/gather
+# tensors at different stages of the forward:
+#
+# ```python
+# _cp_plan = {
+#     "": {
+#         "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+#         "encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
+#         "encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
+#     },
+#     "pos_embed": {
+#         0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+#         1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
+#     },
+#     "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
+# }
+# ```
+#
+# The dictionary is a set of module names mapped to their respective CP plan. The inputs/outputs of layers will be
+# split/gathered according to this at the respective module level. Here, the following happens:
+# - "":
+#     we specify that we want to split the various inputs across the sequence dim in the pre-forward hook (i.e. before
+#     the actual forward logic of the QwenImageTransformer2DModel is run, we will splitthe inputs)
+# - "pos_embed":
+#     we specify that we want to split the outputs of the RoPE layer. Since there are two outputs (imag & text freqs),
+#     we can individually specify how they should be split
+# - "proj_out":
+#     before returning to the user, we gather the entire sequence on each rank in the post-forward hook (after the linear
+#     layer forward has run).
+#
+# ContextParallelInput:
+#     specifies how to split the input tensor in the pre-forward or post-forward hook of the layer it is attached to
+#
+# ContextParallelOutput:
+#     specifies how to gather the input tensor in the post-forward hook in the layer it is attached to
+
+from .cp_plan_registers import (
+    ContextParallelismPlanner,
+    ContextParallelismPlannerRegister,
+)
+from .cp_plan_flux import FluxContextParallelismPlanner
+from .cp_plan_qwen_image import QwenImageContextParallelismPlanner
+from .cp_plan_wan import WanContextParallelismPlanner
+from .cp_plan_ltxvideo import LTXVideoContextParallelismPlanner
+from .cp_plan_hunyuan import HunyuanImageContextParallelismPlanner
+from .cp_plan_hunyuan import HunyuanVideoContextParallelismPlanner
+from .cp_plan_cogvideox import CogVideoXContextParallelismPlanner
+from .cp_plan_cogview import CogView3PlusContextParallelismPlanner
+from .cp_plan_cogview import CogView4ContextParallelismPlanner
+from .cp_plan_cosisid import CosisIDContextParallelismPlanner
+from .cp_plan_chroma import ChromaContextParallelismPlanner
+from .cp_plan_pixart import PixArtContextParallelismPlanner
+from .cp_plan_dit import DiTContextParallelismPlanner
+
+try:
+    import nunchaku  # noqa: F401
+
+    _nunchaku_available = True
+except ImportError:
+    _nunchaku_available = False
+
+if _nunchaku_available:
+    from .cp_plan_nunchaku import (  # noqa: F401
+        NunchakuFluxContextParallelismPlanner,
+    )
+    from .cp_plan_nunchaku import (  # noqa: F401
+        NunchakuQwenImageContextParallelismPlanner,
+    )
+
+
+__all__ = [
+    "ContextParallelismPlanner",
+    "ContextParallelismPlannerRegister",
+    "FluxContextParallelismPlanner",
+    "QwenImageContextParallelismPlanner",
+    "WanContextParallelismPlanner",
+    "LTXVideoContextParallelismPlanner",
+    "HunyuanImageContextParallelismPlanner",
+    "HunyuanVideoContextParallelismPlanner",
+    "CogVideoXContextParallelismPlanner",
+    "CogView3PlusContextParallelismPlanner",
+    "CogView4ContextParallelismPlanner",
+    "CosisIDContextParallelismPlanner",
+    "ChromaContextParallelismPlanner",
+    "PixArtContextParallelismPlanner",
+    "DiTContextParallelismPlanner",
+]
+
+if _nunchaku_available:
+    __all__.extend(
+        [
+            "NunchakuFluxContextParallelismPlanner",
+            "NunchakuQwenImageContextParallelismPlanner",
+        ]
+    )