[Hub] feat: explicitly tag to diffusers when using push_to_hub (#6678)

* feat: explicitly tag to diffusers when using push_to_hub

* remove tags.

* reset repo.

* Apply suggestions from code review

Co-authored-by: Patrick von Platen <[email protected]>

* fix: tests

* fix: push_to_hub behaviour for tagging from save_pretrained

* Apply suggestions from code review

Co-authored-by: Lucain <[email protected]>

* Apply suggestions from code review

Co-authored-by: Lucain <[email protected]>

* import fixes.

* add library name to existing model card.

* add: standalone test for generate_model_card

* fix tests for standalone method

* moved library_name to a better place.

* merge create_model_card and generate_model_card.

* fix test

* address lucain's comments

* fix return identation

* Apply suggestions from code review

Co-authored-by: Lucain <[email protected]>

* address further comments.

* Update src/diffusers/pipelines/pipeline_utils.py

Co-authored-by: Lucain <[email protected]>

---------

Co-authored-by: Patrick von Platen <[email protected]>
Co-authored-by: Lucain <[email protected]>

Author: 3 people

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, load_or_create_model_card, populate_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 = load_or_create_model_card(repo_id, token=token)
+            model_card = populate_model_card(model_card)
+            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 load_or_create_model_card, populate_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 = load_or_create_model_card(repo_id, token=token, is_pipeline=True)
+            model_card = populate_model_card(model_card)
+            model_card.save(os.path.join(save_directory, "README.md"))
+
             self._upload_folder(
                 save_directory,
                 repo_id,

src/diffusers/utils/hub_utils.py

@@ -28,7 +28,6 @@
     ModelCard,
     ModelCardData,
     create_repo,
-    get_full_repo_name,
     hf_hub_download,
     upload_folder,
 )
@@ -67,7 +66,6 @@
 logger = get_logger(__name__)
 
 
-MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "model_card_template.md"
 SESSION_ID = uuid4().hex
 
 
@@ -95,53 +93,45 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str:
     return ua
 
 
-def create_model_card(args, model_name):
+def load_or_create_model_card(
+    repo_id_or_path: Optional[str] = None, token: Optional[str] = None, is_pipeline: bool = False
+) -> ModelCard:
+    """
+    Loads or creates a model card.
+
+    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`].
+    """
     if not is_jinja_available():
         raise ValueError(
             "Modelcard rendering is based on Jinja templates."
             " Please make sure to have `jinja` installed before using `create_model_card`."
             " To install it, please run `pip install Jinja2`."
         )
 
-    if hasattr(args, "local_rank") and args.local_rank not in [-1, 0]:
-        return
-
-    hub_token = args.hub_token if hasattr(args, "hub_token") else None
-    repo_name = get_full_repo_name(model_name, token=hub_token)
-
-    model_card = ModelCard.from_template(
-        card_data=ModelCardData(  # Card metadata object that will be converted to YAML block
-            language="en",
-            license="apache-2.0",
-            library_name="diffusers",
-            tags=[],
-            datasets=args.dataset_name,
-            metrics=[],
-        ),
-        template_path=MODEL_CARD_TEMPLATE_PATH,
-        model_name=model_name,
-        repo_name=repo_name,
-        dataset_name=args.dataset_name if hasattr(args, "dataset_name") else None,
-        learning_rate=args.learning_rate,
-        train_batch_size=args.train_batch_size,
-        eval_batch_size=args.eval_batch_size,
-        gradient_accumulation_steps=(
-            args.gradient_accumulation_steps if hasattr(args, "gradient_accumulation_steps") else None
-        ),
-        adam_beta1=args.adam_beta1 if hasattr(args, "adam_beta1") else None,
-        adam_beta2=args.adam_beta2 if hasattr(args, "adam_beta2") else None,
-        adam_weight_decay=args.adam_weight_decay if hasattr(args, "adam_weight_decay") else None,
-        adam_epsilon=args.adam_epsilon if hasattr(args, "adam_epsilon") else None,
-        lr_scheduler=args.lr_scheduler if hasattr(args, "lr_scheduler") else None,
-        lr_warmup_steps=args.lr_warmup_steps if hasattr(args, "lr_warmup_steps") else None,
-        ema_inv_gamma=args.ema_inv_gamma if hasattr(args, "ema_inv_gamma") else None,
-        ema_power=args.ema_power if hasattr(args, "ema_power") else None,
-        ema_max_decay=args.ema_max_decay if hasattr(args, "ema_max_decay") else None,
-        mixed_precision=args.mixed_precision,
-    )
-
-    card_path = os.path.join(args.output_dir, "README.md")
-    model_card.save(card_path)
+    try:
+        # Check if the model card is present on the remote repo
+        model_card = ModelCard.load(repo_id_or_path, 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)
+
+    return model_card
+
+
+def populate_model_card(model_card: ModelCard) -> ModelCard:
+    """Populates the `model_card` with library name."""
+    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):
@@ -435,6 +425,10 @@ 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 = load_or_create_model_card(repo_id, token=token)
+        model_card = populate_model_card(model_card)
+
         # Save all files.
         save_kwargs = {"safe_serialization": safe_serialization}
         if "Scheduler" not in self.__class__.__name__:
@@ -443,6 +437,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,

src/diffusers/utils/model_card_template.md

@@ -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/models/test_modeling_common.py

@@ -15,37 +15,15 @@
 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 load_or_create_model_card, populate_model_card
 
 
 class CreateModelCardTest(unittest.TestCase):
-    @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"
+    def test_generate_model_card_with_library_name(self):
         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
-
-            # Model card mush be rendered and saved
-            diffusers.utils.hub_utils.create_model_card(args, model_name="model_name")
-            self.assertTrue((Path(tmpdir) / "README.md").is_file())
+            file_path = Path(tmpdir) / "README.md"
+            file_path.write_text("---\nlibrary_name: foo\n---\nContent\n")
+            model_card = load_or_create_model_card(file_path)
+            populate_model_card(model_card)
+            assert model_card.data.library_name == "foo"

tests/others/test_hub_utils.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()`