apply suggestions from review

Author: a-r-r-o-w

src/diffusers/pipelines/pipeline_loading_utils.py

@@ -40,7 +40,6 @@
     is_transformers_available,
     logging,
 )
-from ..utils.hub_utils import _check_legacy_sharding_variant_format
 from ..utils.torch_utils import is_compiled_module
 
 
@@ -841,25 +840,6 @@ def get_connected_passed_kwargs(prefix):
     return init_kwargs
 
 
-def _maybe_raise_warning_for_variant_checkpoint_format(
-    folder: Optional[str] = None, filenames: Optional[List[str]] = None, variant: Optional[str] = None
-) -> None:
-    if variant is not None and _check_legacy_sharding_variant_format(
-        folder=folder, filenames=filenames, variant=variant
-    ):
-        warn_msg = (
-            f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. "
-            "Please check your files carefully:\n\n"
-            "- Correct format example: diffusion_pytorch_model.fp16-00003-of-00003.safetensors\n"
-            "- Deprecated format example: diffusion_pytorch_model-00001-of-00002.fp16.safetensors\n\n"
-            "If you find any files in the deprecated format:\n"
-            "1. Remove all existing checkpoint files for this variant.\n"
-            "2. Re-obtain the correct files by running `save_pretrained()`.\n\n"
-            "This will ensure you're using the most up-to-date and compatible checkpoint format."
-        )
-        logger.warning(warn_msg)
-
-
 def _get_custom_components_and_folders(
     pretrained_model_name: str,
     config_dict: Dict[str, Any],

src/diffusers/pipelines/pipeline_utils.py

@@ -70,6 +70,7 @@
     CONNECTED_PIPES_KEYS,
     CUSTOM_PIPELINE_FILE_NAME,
     LOADABLE_CLASSES,
+    _check_legacy_sharding_variant_format,
     _fetch_class_library_tuple,
     _get_custom_components_and_folders,
     _get_custom_pipeline_class,
@@ -78,7 +79,6 @@
     _get_pipeline_class,
     _identify_model_variants,
     _maybe_raise_warning_for_inpainting,
-    _maybe_raise_warning_for_variant_checkpoint_format,
     _resolve_custom_pipeline_and_cls,
     _unwrap_model,
     _update_init_kwargs_with_connected_pipeline,
@@ -738,7 +738,18 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
 
         # The variant filenames can have the legacy sharding checkpoint format that we check and throw
         # a warning if detected.
-        _maybe_raise_warning_for_variant_checkpoint_format(folder=cached_folder, variant=variant)
+        if variant is not None and _check_legacy_sharding_variant_format(folder=cached_folder, variant=variant):
+            warn_msg = (
+                f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. "
+                "Please check your files carefully:\n\n"
+                "- Correct format example: diffusion_pytorch_model.fp16-00003-of-00003.safetensors\n"
+                "- Deprecated format example: diffusion_pytorch_model-00001-of-00002.fp16.safetensors\n\n"
+                "If you find any files in the deprecated format:\n"
+                "1. Remove all existing checkpoint files for this variant.\n"
+                "2. Re-obtain the correct files by running `save_pretrained()`.\n\n"
+                "This will ensure you're using the most up-to-date and compatible checkpoint format."
+            )
+            logger.warning(warn_msg)
 
         config_dict = cls.load_config(cached_folder)
 
@@ -1261,7 +1272,18 @@ def download(cls, pretrained_model_name, **kwargs) -> Union[str, os.PathLike]:
 
         if not local_files_only:
             filenames = {sibling.rfilename for sibling in info.siblings}
-            _maybe_raise_warning_for_variant_checkpoint_format(filenames=filenames, variant=variant)
+            if variant is not None and _check_legacy_sharding_variant_format(filenames=filenames, variant=variant):
+                warn_msg = (
+                    f"Warning: The repository contains sharded checkpoints for variant '{variant}' maybe in a deprecated format. "
+                    "Please check your files carefully:\n\n"
+                    "- Correct format example: diffusion_pytorch_model.fp16-00003-of-00003.safetensors\n"
+                    "- Deprecated format example: diffusion_pytorch_model-00001-of-00002.fp16.safetensors\n\n"
+                    "If you find any files in the deprecated format:\n"
+                    "1. Remove all existing checkpoint files for this variant.\n"
+                    "2. Re-obtain the correct files by running `save_pretrained()`.\n\n"
+                    "This will ensure you're using the most up-to-date and compatible checkpoint format."
+                )
+                logger.warning(warn_msg)
 
             model_filenames, variant_filenames = variant_compatible_siblings(filenames, variant=variant)
 

tests/pipelines/test_pipelines.py

@@ -578,7 +578,7 @@ def test_download_variants_with_sharded_checkpoints(self):
 
     def test_download_legacy_variants_with_sharded_ckpts_raises_warning(self):
         repo_id = "hf-internal-testing/tiny-stable-diffusion-pipe-variants-all-kinds"
-        logger = logging.get_logger("diffusers.pipelines.pipeline_loading_utils")
+        logger = logging.get_logger("diffusers.pipelines.pipeline_utils")
         deprecated_warning_msg = "Warning: The repository contains sharded checkpoints for variant"
 
         for is_local in [True, False]: