[Hub] feat: explicitly tag to diffusers when using push_to_hub

src/diffusers/models/modeling_utils.py

@@ -41,7 +41,7 @@
     is_torch_version,
     logging,
 )
-from ..utils.hub_utils import PushToHubMixin
+from ..utils.hub_utils import PushToHubMixin, create_and_tag_model_card
 
 
 logger = logging.get_logger(__name__)
@@ -375,6 +375,11 @@ def save_pretrained(
         logger.info(f"Model weights saved in {os.path.join(save_directory, weights_name)}")
 
         if push_to_hub:
+            # Create a new empty model card and eventually tag it
+            model_card = create_and_tag_model_card(repo_id, token=token)
+            # Update model card if needed:
+            model_card.save(os.path.join(save_directory, "README.md"))
+
             self._upload_folder(
                 save_directory,
                 repo_id,

src/diffusers/pipelines/pipeline_utils.py

@@ -60,6 +60,7 @@
     logging,
     numpy_to_pil,
 )
+from ..utils.hub_utils import create_and_tag_model_card
 from ..utils.torch_utils import is_compiled_module
 
 
@@ -720,6 +721,11 @@ def is_saveable_module(name, value):
         self.save_config(save_directory)
 
         if push_to_hub:
+            # Create a new empty model card and eventually tag it
+            model_card = create_and_tag_model_card(repo_id, token=token, is_pipeline=True)
+            # Update model card if needed:
+            model_card.save(os.path.join(save_directory, "README.md"))
+
             self._upload_folder(
                 save_directory,
                 repo_id,

src/diffusers/utils/hub_utils.py

@@ -144,6 +144,32 @@ def create_model_card(args, model_name):
     model_card.save(card_path)
 
 
+# Taken from `transformers`
+def create_and_tag_model_card(repo_id: str, token: Optional[str] = None, is_pipeline=False):
+    """
+    Creates or loads an existing model card and tags it with the `library_name`.
+
+    Args:
+        repo_id (`str`):
+            The repo_id where to look for the model card.
+        token (`str`, *optional*):
+            Authentication token, obtained with `huggingface_hub.HfApi.login` method. Will default to the stored token.
+        is_pipeline (`bool`, *optional*):
+            Boolean to indicate if we're adding tag to a [`DiffusionPipeline`].
+    """
+    try:
+        # Check if the model card is present on the remote repo
+        model_card = ModelCard.load(repo_id, token=token, ignore_metadata_errors=False)
+    except EntryNotFoundError:
+        # Otherwise create a simple model card from template
+        component = "pipeline" if is_pipeline else "model"
+        model_description = f"This is the model card of a 🧨 diffusers {component} that has been pushed on the Hub. This model card has been automatically generated."
+        card_data = ModelCardData(library_name="diffusers")
+        model_card = ModelCard.from_template(card_data, model_description=model_description)
+
+    return model_card
+
+
 def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str] = None):
     """
     Extracts the commit hash from a resolved filename toward a cache file.
@@ -435,6 +461,9 @@ def push_to_hub(
         """
         repo_id = create_repo(repo_id, private=private, token=token, exist_ok=True).repo_id
 
+        # Create a new empty model card and eventually tag it
+        model_card = create_and_tag_model_card(repo_id, token=token)
+
         # Save all files.
         save_kwargs = {"safe_serialization": safe_serialization}
         if "Scheduler" not in self.__class__.__name__:
@@ -443,6 +472,9 @@ def push_to_hub(
         with tempfile.TemporaryDirectory() as tmpdir:
             self.save_pretrained(tmpdir, **save_kwargs)
 
+            # Update model card if needed:
+            model_card.save(os.path.join(tmpdir, "README.md"))
+
             return self._upload_folder(
                 tmpdir,
                 repo_id,

tests/models/test_modeling_common.py

@@ -24,7 +24,8 @@
 import numpy as np
 import requests_mock
 import torch
-from huggingface_hub import delete_repo
+from huggingface_hub import ModelCard, delete_repo
+from huggingface_hub.utils import is_jinja_available
 from requests.exceptions import HTTPError
 
 from diffusers.models import UNet2DConditionModel
@@ -732,3 +733,26 @@ def test_push_to_hub_in_organization(self):
 
         # Reset repo
         delete_repo(self.org_repo_id, token=TOKEN)
+
+    @unittest.skipIf(
+        not is_jinja_available(),
+        reason="Model card tests cannot be performed with Jinja installed.",
+    )
+    def test_push_to_hub_library_name(self):
+        model = UNet2DConditionModel(
+            block_out_channels=(32, 64),
+            layers_per_block=2,
+            sample_size=32,
+            in_channels=4,
+            out_channels=4,
+            down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
+            up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
+            cross_attention_dim=32,
+        )
+        model.push_to_hub(self.repo_id, token=TOKEN)
+
+        model_card = ModelCard.load(f"{USER}/{self.repo_id}", token=TOKEN).data
+        assert model_card.library_name == "diffusers"
+
+        # Reset repo
+        delete_repo(self.repo_id, token=TOKEN)

tests/pipelines/test_pipelines_common.py

@@ -13,7 +13,8 @@
 import numpy as np
 import PIL.Image
 import torch
-from huggingface_hub import delete_repo
+from huggingface_hub import ModelCard, delete_repo
+from huggingface_hub.utils import is_jinja_available
 from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
 
 import diffusers
@@ -1142,6 +1143,21 @@ def test_push_to_hub_in_organization(self):
         # Reset repo
         delete_repo(self.org_repo_id, token=TOKEN)
 
+    @unittest.skipIf(
+        not is_jinja_available(),
+        reason="Model card tests cannot be performed with Jinja installed.",
+    )
+    def test_push_to_hub_library_name(self):
+        components = self.get_pipeline_components()
+        pipeline = StableDiffusionPipeline(**components)
+        pipeline.push_to_hub(self.repo_id, token=TOKEN)
+
+        model_card = ModelCard.load(f"{USER}/{self.repo_id}", token=TOKEN).data
+        assert model_card.library_name == "diffusers"
+
+        # Reset repo
+        delete_repo(self.repo_id, token=TOKEN)
+
 
 # For SDXL and its derivative pipelines (such as ControlNet), we have the text encoders
 # and the tokenizers as optional components. So, we need to override the `test_save_load_optional_components()`