review feedback.

Author: sayakpaul

src/diffusers/models/model_loading_utils.py

@@ -407,8 +407,7 @@ def _load_shard_files_with_threadpool(
     low_cpu_mem_usage=False,
 ):
     # Do not spawn anymore workers than you need
-    num_workers = int(os.environ.get("HF_PARALLEL_LOADING_WORKERS", str(DEFAULT_HF_PARALLEL_LOADING_WORKERS)))
-    num_workers = min(len(shard_files), num_workers)
+    num_workers = min(len(shard_files), DEFAULT_HF_PARALLEL_LOADING_WORKERS)
 
     logger.info(f"Loading model weights in parallel with {num_workers} workers...")
 

src/diffusers/models/modeling_utils.py

@@ -44,6 +44,7 @@
     CONFIG_NAME,
     ENV_VARS_TRUE_VALUES,
     FLAX_WEIGHTS_NAME,
+    HF_PARALLEL_LOADING_FLAG,
     SAFE_WEIGHTS_INDEX_NAME,
     SAFETENSORS_WEIGHTS_NAME,
     WEIGHTS_INDEX_NAME,
@@ -961,7 +962,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         dduf_entries: Optional[Dict[str, DDUFEntry]] = kwargs.pop("dduf_entries", None)
         disable_mmap = kwargs.pop("disable_mmap", False)
 
-        is_parallel_loading_enabled = os.environ.get("HF_ENABLE_PARALLEL_LOADING", "").upper() in ENV_VARS_TRUE_VALUES
+        is_parallel_loading_enabled = os.environ.get(HF_PARALLEL_LOADING_FLAG, "").upper() in ENV_VARS_TRUE_VALUES
         if is_parallel_loading_enabled and not low_cpu_mem_usage:
             raise NotImplementedError("Parallel loading is not supported when not using `low_cpu_mem_usage`.")
 

src/diffusers/utils/__init__.py

@@ -26,6 +26,7 @@
     FLAX_WEIGHTS_NAME,
     GGUF_FILE_EXTENSION,
     HF_MODULES_CACHE,
+    HF_PARALLEL_LOADING_FLAG,
     HUGGINGFACE_CO_RESOLVE_ENDPOINT,
     MIN_PEFT_VERSION,
     ONNX_EXTERNAL_WEIGHTS_NAME,

src/diffusers/utils/constants.py

@@ -44,6 +44,7 @@
 DIFFUSERS_ATTN_BACKEND = os.getenv("DIFFUSERS_ATTN_BACKEND", "native")
 DIFFUSERS_ATTN_CHECKS = os.getenv("DIFFUSERS_ATTN_CHECKS", "0") in ENV_VARS_TRUE_VALUES
 DEFAULT_HF_PARALLEL_LOADING_WORKERS = 8
+HF_PARALLEL_LOADING_FLAG = "HF_ENABLE_PARALLEL_LOADING"
 
 # Below should be `True` if the current version of `peft` and `transformers` are compatible with
 # PEFT backend. Will automatically fall back to PEFT backend if the correct versions of the libraries are