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

src/diffusers/models/modeling_utils.py

@@ -42,7 +42,7 @@
     is_torch_version,
     logging,
 )
-from ..utils.hub_utils import PushToHubMixin
+from ..utils.hub_utils import PushToHubMixin, generate_model_card
 
 
 logger = logging.get_logger(__name__)
@@ -377,6 +377,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 = generate_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 generate_model_card
 from ..utils.torch_utils import is_compiled_module
 
 
@@ -725,6 +726,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 = generate_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,35 @@ def create_model_card(args, model_name):
     model_card.save(card_path)
 
 
+# Taken from `transformers`
+def generate_model_card(repo_id: str, token: Optional[str] = None, is_pipeline: bool = False) -> ModelCard:
+    """
+    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. Will default to the stored token. See https://huggingface.co/settings/token for more details.
+        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)
+    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()
+        model_card = ModelCard.from_template(card_data, model_description=model_description)
+
+    if model_card.data.library_name is None:
+        model_card.data.library_name = "diffusers"
+
+    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 +464,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 = generate_model_card(repo_id, token=token)
+
         # Save all files.
         save_kwargs = {"safe_serialization": safe_serialization}
         if "Scheduler" not in self.__class__.__name__:
@@ -443,6 +475,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 without 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/others/test_hub_utils.py

@@ -12,40 +12,61 @@
 # 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.
+import os
 import unittest
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest.mock import Mock, patch
 
-import diffusers.utils.hub_utils
+from diffusers.utils.hub_utils import create_model_card, generate_model_card
 
 
 class CreateModelCardTest(unittest.TestCase):
+    def create_dummy_args(self, output_dir):
+        # Dummy args values
+        args = Mock()
+        args.output_dir = output_dir
+        args.local_rank = 0
+        args.hub_token = "hub_token"
+        args.dataset_name = "dataset_name"
+        args.learning_rate = 0.01
+        args.train_batch_size = 100000
+        args.eval_batch_size = 10000
+        args.gradient_accumulation_steps = 0.01
+        args.adam_beta1 = 0.02
+        args.adam_beta2 = 0.03
+        args.adam_weight_decay = 0.0005
+        args.adam_epsilon = 0.000001
+        args.lr_scheduler = 1
+        args.lr_warmup_steps = 10
+        args.ema_inv_gamma = 0.001
+        args.ema_power = 0.1
+        args.ema_max_decay = 0.2
+        args.mixed_precision = True
+        return args
+
     @patch("diffusers.utils.hub_utils.get_full_repo_name")
     def test_create_model_card(self, repo_name_mock: Mock) -> None:
         repo_name_mock.return_value = "full_repo_name"
         with TemporaryDirectory() as tmpdir:
-            # Dummy args values
-            args = Mock()
-            args.output_dir = tmpdir
-            args.local_rank = 0
-            args.hub_token = "hub_token"
-            args.dataset_name = "dataset_name"
-            args.learning_rate = 0.01
-            args.train_batch_size = 100000
-            args.eval_batch_size = 10000
-            args.gradient_accumulation_steps = 0.01
-            args.adam_beta1 = 0.02
-            args.adam_beta2 = 0.03
-            args.adam_weight_decay = 0.0005
-            args.adam_epsilon = 0.000001
-            args.lr_scheduler = 1
-            args.lr_warmup_steps = 10
-            args.ema_inv_gamma = 0.001
-            args.ema_power = 0.1
-            args.ema_max_decay = 0.2
-            args.mixed_precision = True
+            args = self.create_dummy_args(output_dir=tmpdir)
 
             # Model card mush be rendered and saved
-            diffusers.utils.hub_utils.create_model_card(args, model_name="model_name")
+            create_model_card(args, model_name="model_name")
             self.assertTrue((Path(tmpdir) / "README.md").is_file())
+
+    def test_generate_existing_model_card_with_library_name(self):
+        with TemporaryDirectory() as tmpdir:
+            content = "hello"
+            with open(os.path.join(tmpdir, "README.md"), "w") as f:
+                f.write(content)
+
+            model_card = generate_model_card(tmpdir)
+            assert model_card.data.library_name == "diffusers"
+
+    def test_generate_model_card_with_library_name(self):
+        with TemporaryDirectory() as tmpdir:
+            model_card = generate_model_card(tmpdir)
+
+        model_card = generate_model_card(tmpdir)
+        assert model_card.data.library_name == "diffusers"

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 without 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()`