Repository: huggingface/diffusers
File: src/diffusers/models/modeling_utils.py (Python)

Code before:
        logger.warning(
            "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
        )

        if not torch.distributed.is_available() and not torch.distributed.is_initialized():
            raise RuntimeError(
                "torch.distributed must be available and initialized before calling `enable_parallelism`."
            )

        from ..hooks.context_parallel import apply_context_parallel
        from .attention import AttentionModuleMixin
        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
        from .attention_processor import Attention, MochiAttention

        if isinstance(config, ContextParallelConfig):
            config = ParallelConfig(context_parallel_config=config)

        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()
        device_type = torch._C._get_accelerator().type
        device_module = torch.get_device_module(device_type)
        device = torch.device(device_type, rank % device_module.device_count())

        attention_classes = (Attention, MochiAttention, AttentionModuleMixin)

        # Step 1: Validate attention backend supports context parallelism if enabled
        if config.context_parallel_config is not None:
            for module in self.modules():
                if not isinstance(module, attention_classes):
                    continue

                processor = module.processor
                if processor is None or not hasattr(processor, "_attention_backend"):
                    continue

                attention_backend = processor._attention_backend
                if attention_backend is None:
                    attention_backend, _ = _AttentionBackendRegistry.get_active_backend()
                else:
                    attention_backend = AttentionBackendName(attention_backend)

                if not _AttentionBackendRegistry._is_context_parallel_available(attention_backend):
                    compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)
                    raise ValueError(
                        f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' "
                        f"is using backend '{attention_backend.value}' which does not support context parallelism. "
                        f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before "
                        f"calling `enable_parallelism()`."
                    )

                # All modules use the same attention processor and backend. We don't need to

Review comment (refactor):
(nit): Seems like we are not documenting other steps. So, maybe we can remove this comment or add other steps as comments.

Code after:
        logger.warning(
            "`enable_parallelism` is an experimental feature. The API may change in the future and breaking changes may be introduced at any time without warning."
        )

        if not torch.distributed.is_available() and not torch.distributed.is_initialized():
            raise RuntimeError(
                "torch.distributed must be available and initialized before calling `enable_parallelism`."
            )

        from ..hooks.context_parallel import apply_context_parallel
        from .attention import AttentionModuleMixin
        from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry
        from .attention_processor import Attention, MochiAttention

        if isinstance(config, ContextParallelConfig):
            config = ParallelConfig(context_parallel_config=config)

        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()
        device_type = torch._C._get_accelerator().type
        device_module = torch.get_device_module(device_type)
        device = torch.device(device_type, rank % device_module.device_count())

        attention_classes = (Attention, MochiAttention, AttentionModuleMixin)

        if config.context_parallel_config is not None:
            for module in self.modules():
                if not isinstance(module, attention_classes):
                    continue

                processor = module.processor
                if processor is None or not hasattr(processor, "_attention_backend"):
                    continue

                attention_backend = processor._attention_backend
                if attention_backend is None:
                    attention_backend, _ = _AttentionBackendRegistry.get_active_backend()
                else:
                    attention_backend = AttentionBackendName(attention_backend)

                if not _AttentionBackendRegistry._is_context_parallel_available(attention_backend):
                    compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)
                    raise ValueError(
                        f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' "
                        f"is using backend '{attention_backend.value}' which does not support context parallelism. "
                        f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before "
                        f"calling `enable_parallelism()`."
                    )

                # All modules use the same attention processor and backend. We don't need to
                # iterate over all modules after checking the first processor