up

Author: sayakpaul

src/diffusers/models/model_loading_utils.py

@@ -47,6 +47,7 @@
     is_torch_version,
     logging,
 )
+from ..utils.distributed_utils import is_torch_dist_rank_zero
 
 
 logger = logging.get_logger(__name__)
@@ -429,8 +430,12 @@ def _load_shard_files_with_threadpool(
         low_cpu_mem_usage=low_cpu_mem_usage,
     )
 
+    tqdm_kwargs = {"total": len(shard_files), "desc": "Loading checkpoint shards"}
+    if not is_torch_dist_rank_zero():
+        tqdm_kwargs["disable"] = True
+
     with ThreadPoolExecutor(max_workers=num_workers) as executor:
-        with logging.tqdm(total=len(shard_files), desc="Loading checkpoint shards") as pbar:
+        with logging.tqdm(**tqdm_kwargs) as pbar:
             futures = [executor.submit(load_one, shard_file) for shard_file in shard_files]
             for future in as_completed(futures):
                 result = future.result()

src/diffusers/models/modeling_utils.py

@@ -59,12 +59,9 @@
     is_torch_version,
     logging,
 )
-from ..utils.hub_utils import (
-    PushToHubMixin,
-    load_or_create_model_card,
-    populate_model_card,
-)
-from ..utils.torch_utils import empty_device_cache, is_torch_dist_rank_zero
+from ..utils.distributed_utils import is_torch_dist_rank_zero
+from ..utils.hub_utils import PushToHubMixin, load_or_create_model_card, populate_model_card
+from ..utils.torch_utils import empty_device_cache
 from ._modeling_parallel import ContextParallelConfig, ContextParallelModelPlan, ParallelConfig
 from .model_loading_utils import (
     _caching_allocator_warmup,

src/diffusers/pipelines/pipeline_utils.py

@@ -67,8 +67,9 @@
     logging,
     numpy_to_pil,
 )
+from ..utils.distributed_utils import is_torch_dist_rank_zero
 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, is_torch_dist_rank_zero
+from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module
 
 
 if is_torch_npu_available():

src/diffusers/utils/distributed_utils.py

@@ -0,0 +1,36 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+def is_torch_dist_rank_zero() -> bool:
+    if torch is None:
+        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

src/diffusers/utils/logging.py

@@ -32,6 +32,8 @@
 
 from tqdm import auto as tqdm_lib
 
+from .distributed_utils import is_torch_dist_rank_zero
+
 
 _lock = threading.Lock()
 _default_handler: Optional[logging.Handler] = None
@@ -47,6 +49,22 @@
 _default_log_level = logging.WARNING
 
 _tqdm_active = True
+_rank_zero_filter = None
+
+
+class _RankZeroFilter(logging.Filter):
+    def filter(self, record):
+        return is_torch_dist_rank_zero()
+
+
+def _ensure_rank_zero_filter(logger: logging.Logger) -> None:
+    global _rank_zero_filter
+
+    if _rank_zero_filter is None:
+        _rank_zero_filter = _RankZeroFilter()
+
+    if not any(isinstance(f, _RankZeroFilter) for f in logger.filters):
+        logger.addFilter(_rank_zero_filter)
 
 
 def _get_default_logging_level() -> int:
@@ -90,6 +108,7 @@ def _configure_library_root_logger() -> None:
         library_root_logger.addHandler(_default_handler)
         library_root_logger.setLevel(_get_default_logging_level())
         library_root_logger.propagate = False
+        _ensure_rank_zero_filter(library_root_logger)
 
 
 def _reset_library_root_logger() -> None:
@@ -120,7 +139,9 @@ def get_logger(name: Optional[str] = None) -> logging.Logger:
         name = _get_library_name()
 
     _configure_library_root_logger()
-    return logging.getLogger(name)
+    logger = logging.getLogger(name)
+    _ensure_rank_zero_filter(logger)
+    return logger
 
 
 def get_verbosity() -> int:

src/diffusers/utils/torch_utils.py

@@ -143,23 +143,6 @@ 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,