up

Author: sayakpaul

src/diffusers/models/modeling_utils.py

@@ -64,7 +64,7 @@
     load_or_create_model_card,
     populate_model_card,
 )
-from ..utils.torch_utils import empty_device_cache
+from ..utils.torch_utils import empty_device_cache, is_torch_dist_rank_zero
 from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig
 from .model_loading_utils import (
     _caching_allocator_warmup,
@@ -1672,7 +1672,10 @@ def _load_pretrained_model(
         else:
             shard_files = resolved_model_file
             if len(resolved_model_file) > 1:
-                shard_files = logging.tqdm(resolved_model_file, desc="Loading checkpoint shards")
+                shard_tqdm_kwargs = {"desc": "Loading checkpoint shards"}
+                if not is_torch_dist_rank_zero():
+                    shard_tqdm_kwargs["disable"] = True
+                shard_files = logging.tqdm(resolved_model_file, **shard_tqdm_kwargs)
 
             for shard_file in shard_files:
                 offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file)

src/diffusers/pipelines/pipeline_utils.py

@@ -68,7 +68,7 @@
     numpy_to_pil,
 )
 from ..utils.hub_utils import _check_legacy_sharding_variant_format, load_or_create_model_card, populate_model_card
-from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module
+from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module, is_torch_dist_rank_zero
 
 
 if is_torch_npu_available():
@@ -983,7 +983,7 @@ def load_module(name, value):
         current_device_map = None
         _maybe_warn_for_wrong_component_in_quant_config(init_dict, quantization_config)
         logging_tqdm_kwargs = {"desc": "Loading pipeline components..."}
-        if cls._progress_bar_disabled_for_rank():
+        if not is_torch_dist_rank_zero():
             logging_tqdm_kwargs["disable"] = True
 
         for name, (library_name, class_name) in logging.tqdm(init_dict.items(), **logging_tqdm_kwargs):
@@ -1914,7 +1914,7 @@ def progress_bar(self, iterable=None, total=None):
 
         progress_bar_config = dict(self._progress_bar_config)
         if "disable" not in progress_bar_config:
-            progress_bar_config["disable"] = self._progress_bar_disabled_for_rank()
+            progress_bar_config["disable"] = not is_torch_dist_rank_zero()
 
         if iterable is not None:
             return tqdm(iterable, **progress_bar_config)
@@ -1926,15 +1926,6 @@ def progress_bar(self, iterable=None, total=None):
     def set_progress_bar_config(self, **kwargs):
         self._progress_bar_config = kwargs
 
-    @staticmethod
-    def _progress_bar_disabled_for_rank():
-        if torch.distributed.is_available() and torch.distributed.is_initialized():
-            try:
-                return torch.distributed.get_rank() != 0
-            except (RuntimeError, ValueError):
-                return False
-        return False
-
     def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
         r"""
         Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this

src/diffusers/utils/torch_utils.py

@@ -143,6 +143,23 @@ def backend_supports_training(device: str):
     return BACKEND_SUPPORTS_TRAINING[device]
 
 
+def is_torch_dist_rank_zero() -> bool:
+    if not is_torch_available():
+        return True
+
+    dist_module = getattr(torch, "distributed", None)
+    if dist_module is None or not dist_module.is_available():
+        return True
+
+    if not dist_module.is_initialized():
+        return True
+
+    try:
+        return dist_module.get_rank() == 0
+    except (RuntimeError, ValueError):
+        return True
+
+
 def randn_tensor(
     shape: Union[Tuple, List],
     generator: Optional[Union[List["torch.Generator"], "torch.Generator"]] = None,