Diffusers crash due to hard-dependency on `torch._C._distributed_c10d`

[vladmandic]

Describe the bug

BriaTransformer2DModel introduces a hard dependency on torch._C._distributed_c10d which is NOT present in all torch variants/flavors.

Any torch._C import should be guarded with if-available, not a hard dependency on the entire library.
Without that, diffusers stop being a cross-platform solution.

Reproduction

Attempt to import diffusers, it will immediately fail
Downstream issue: vladmandic/sdnext#4240

Logs

│E:\Work\LLM\sdnext-amd\venv\Lib\site-packages\diffusers\models\transformers\__init__.py:20 in <module>                │
│                                                                                                                      │
│  19     from .transformer_allegro import AllegroTransformer3DModel                                                   │
│❱ 20     from .transformer_bria import BriaTransformer2DModel                                                         │
│  21     from .transformer_chroma import ChromaTransformer2DModel                                                     │
│                                                                                                                      │
│E:\Work\LLM\sdnext-amd\venv\Lib\site-packages\diffusers\models\transformers\transformer_bria.py:14 in <module>        │
│                                                                                                                      │
│   13 from ..attention import AttentionModuleMixin, FeedForward                                                       │
│❱  14 from ..attention_dispatch import dispatch_attention_fn                                                          │
│   15 from ..cache_utils import CacheMixin                                                                            │
│                                                                                                                      │
│E:\Work\LLM\sdnext-amd\venv\Lib\site-packages\diffusers\models\attention_dispatch.py:23 in <module>                   │
│                                                                                                                      │
│    22 import torch                                                                                                   │
│❱   23 import torch.distributed._functional_collectives as funcol                                                     │
│    24                                                                                                                │
│                                                                                                                      │
│E:\Work\LLM\sdnext-amd\venv\Lib\site-packages\torch\distributed\_functional_collectives.py:9 in <module>              │
│                                                                                                                      │
│     8 import torch.distributed as dist                                                                               │
│❱    9 import torch.distributed.distributed_c10d as c10d                                                              │
│    10 from torch.distributed.device_mesh import DeviceMesh                                                           │
│                                                                                                                      │
│E:\Work\LLM\sdnext-amd\venv\Lib\site-packages\torch\distributed\distributed_c10d.py:23 in <module>                    │
│                                                                                                                      │
│    22 from torch._C import _DistStoreError as DistStoreError                                                         │
│❱   23 from torch._C._distributed_c10d import (                                                                       │
│    24     _DistributedBackendOptions,                                                                                │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

System Info

diffusers main branch: 64a5187
torch-rocm for windows preview build, specifically from https://repo.radeon.com/rocm/windows/rocm-rel-6.4.4/torch/

Who can help?

@yiyixuxu @sayakpaul @DN6

[0xDELUXA]

We had the same issue with TheRock wheels on Windows, where transformers tried to import this exact torch._C._distributed_c10d package.
The solution came with huggingface/transformers#40038 which checks whether it is available. If not, it does not try to import it, and no error is raised. I think we need something like this here too.

[vladmandic]

in diffusers, the issue is of unsafe imports is introduced via #11941 which adds optional context-parallel features.
its not specific to bria transformer specifically, its just where this import sequence is first triggered.

[0xDELUXA]

I see. Looking forward to a fix.

[sayakpaul]

#12420

[vladmandic]

issue is NOT resolved as there are other code paths in diffuers that lead to same problem. ANY torch._C import should be safe-guarded.

│E:\AI\sdnext-rock\modules\modular_guiders.py:9 in <module>                                                                                                                                                     │
│                                                                                                                                                                                                               │
│   8     'Default': { 'cls': None, 'args': {}, },                                                                                                                                                              │
│❱  9     'CFG: ClassifierFreeGuidance': { 'cls': diffusers.ClassifierFreeGuidance, 'args': {} },                                                                                                               │
│  10     'Auto: AutoGuidance': { 'cls': diffusers.AutoGuidance, 'args': { 'dropout': 1.0, 'auto_guidance_layers': [7, 8, 9], 'auto_guidance_config': None } },                                                 │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\utils\import_utils.py:964 in __getattr__                                                                                                                    │
│                                                                                                                                                                                                               │
│  963         elif name in self._class_to_module.keys():                                                                                                                                                       │
│❱ 964             module = self._get_module(self._class_to_module[name])                                                                                                                                       │
│  965             value = getattr(module, name)                                                                                                                                                                │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\utils\import_utils.py:976 in _get_module                                                                                                                    │
│                                                                                                                                                                                                               │
│  975         except Exception as e:                                                                                                                                                                           │
│❱ 976             raise RuntimeError(                                                                                                                                                                          │
│  977                 f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its"                                                                                     │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Failed to import diffusers.guiders because of the following error (look up to see its traceback):
No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

[vladmandic]

and one more code path causing same crash!

│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\loaders\single_file_model.py:485 in from_single_file                                                                                                        │
│                                                                                                                                                                                                               │
│  484 │   │   if torch_dtype is not None and hf_quantizer is None:                                                                                                                                             │
│❱ 485 │   │   │   model.to(torch_dtype)                                                                                                                                                                        │
│  486                                                                                                                                                                                                          │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\models\modeling_utils.py:1377 in to                                                                                                                         │
│                                                                                                                                                                                                               │
│  1376 │   def to(self, *args, **kwargs):                                                                                                                                                                      │
│❱ 1377 │   │   from ..hooks.group_offloading import _is_group_offload_enabled                                                                                                                                  │
│  1378                                                                                                                                                                                                         │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\hooks\__init__.py:19 in <module>                                                                                                                            │
│                                                                                                                                                                                                               │
│  18 if is_torch_available():                                                                                                                                                                                  │
│❱ 19 │   from .context_parallel import apply_context_parallel                                                                                                                                                  │
│  20 │   from .faster_cache import FasterCacheConfig, apply_faster_cache                                                                                                                                       │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\diffusers\hooks\context_parallel.py:20 in <module>                                                                                                                    │
│                                                                                                                                                                                                               │
│   19 import torch                                                                                                                                                                                             │
│❱  20 import torch.distributed._functional_collectives as funcol                                                                                                                                               │
│   21                                                                                                                                                                                                          │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\torch\distributed\_functional_collectives.py:9 in <module>                                                                                                            │
│                                                                                                                                                                                                               │
│     8 import torch.distributed as dist                                                                                                                                                                        │
│❱    9 import torch.distributed.distributed_c10d as c10d                                                                                                                                                       │
│    10 from torch.distributed.device_mesh import DeviceMesh                                                                                                                                                    │
│                                                                                                                                                                                                               │
│E:\AI\sdnext-rock\venv\Lib\site-packages\torch\distributed\distributed_c10d.py:24 in <module>                                                                                                                  │
│                                                                                                                                                                                                               │
│    23 from torch._C import _DistStoreError as DistStoreError                                                                                                                                                  │
│❱   24 from torch._C._distributed_c10d import (                                                                                                                                                                │
│    25 │   _DistributedBackendOptions,                                                                                                                                                                         │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package

[sayakpaul]

Sorry
#12425