apply review suggestions

Author: a-r-r-o-w

src/diffusers/__init__.py

@@ -261,7 +261,6 @@
             "WanTransformer3DModel",
             "WanVACETransformer3DModel",
             "attention_backend",
-            "enable_parallelism",
         ]
     )
     _import_structure["modular_pipelines"].extend(
@@ -942,7 +941,6 @@
             WanTransformer3DModel,
             WanVACETransformer3DModel,
             attention_backend,
-            enable_parallelism,
         )
         from .modular_pipelines import ComponentsManager, ComponentSpec, ModularPipeline, ModularPipelineBlocks
         from .optimization import (

src/diffusers/models/__init__.py

@@ -25,7 +25,7 @@
 _import_structure = {}
 
 if is_torch_available():
-    _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig", "enable_parallelism"]
+    _import_structure["_modeling_parallel"] = ["ContextParallelConfig", "ParallelConfig"]
     _import_structure["adapter"] = ["MultiAdapter", "T2IAdapter"]
     _import_structure["attention_dispatch"] = ["AttentionBackendName", "attention_backend"]
     _import_structure["auto_model"] = ["AutoModel"]
@@ -120,7 +120,7 @@
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     if is_torch_available():
-        from ._modeling_parallel import ContextParallelConfig, ParallelConfig, enable_parallelism
+        from ._modeling_parallel import ContextParallelConfig, ParallelConfig
         from .adapter import MultiAdapter, T2IAdapter
         from .attention_dispatch import AttentionBackendName, attention_backend
         from .auto_model import AutoModel

src/diffusers/models/_modeling_parallel.py

@@ -241,62 +241,3 @@ def __repr__(self):
 #
 # ContextParallelOutput:
 #     specifies how to gather the input tensor in the post-forward hook in the layer it is attached to
-
-_ENABLE_PARALLELISM_WARN_ONCE = False
-
-
-@contextlib.contextmanager
-def enable_parallelism(model_or_pipeline: Union["DiffusionPipeline", "ModelMixin"]):
-    """
-    A context manager to set the parallelism context for models or pipelines that have been parallelized.
-
-    Args:
-        model_or_pipeline (`DiffusionPipeline` or `ModelMixin`):
-            The model or pipeline to set the parallelism context for. The model or pipeline must have been parallelized
-            with `.enable_parallelism(ParallelConfig(...), ...)` before using this context manager.
-    """
-
-    from diffusers import DiffusionPipeline, ModelMixin
-
-    from .attention_dispatch import _AttentionBackendRegistry
-
-    global _ENABLE_PARALLELISM_WARN_ONCE
-    if not _ENABLE_PARALLELISM_WARN_ONCE:
-        logger.warning(
-            "Support for `enable_parallelism` is experimental and the API may be subject to change in the future."
-        )
-        _ENABLE_PARALLELISM_WARN_ONCE = True
-
-    if isinstance(model_or_pipeline, DiffusionPipeline):
-        parallelized_components = [
-            (name, component)
-            for name, component in model_or_pipeline.components.items()
-            if getattr(component, "_parallel_config", None) is not None
-        ]
-        if len(parallelized_components) > 1:
-            raise ValueError(
-                "Enabling parallelism on a pipeline is not possible when multiple internal components are parallelized. Please run "
-                "different stages of the pipeline separately with `enable_parallelism` on each component manually."
-            )
-        if len(parallelized_components) == 0:
-            raise ValueError(
-                "No parallelized components found in the pipeline. Please ensure at least one component is parallelized."
-            )
-        _, model_or_pipeline = parallelized_components[0]
-    elif isinstance(model_or_pipeline, ModelMixin):
-        if getattr(model_or_pipeline, "_parallel_config", None) is None:
-            raise ValueError(
-                "The model is not parallelized. Please ensure the model is parallelized with `.parallelize()` before using this context manager."
-            )
-    else:
-        raise TypeError(
-            f"Expected a `DiffusionPipeline` or `ModelMixin` instance, but got {type(model_or_pipeline)}. Please provide a valid model or pipeline."
-        )
-
-    # TODO: needs to be updated when more parallelism strategies are supported
-    old_parallel_config = _AttentionBackendRegistry._parallel_config
-    _AttentionBackendRegistry._parallel_config = model_or_pipeline._parallel_config.context_parallel_config
-
-    yield
-
-    _AttentionBackendRegistry._parallel_config = old_parallel_config