Standardize model cards

examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py

@@ -66,6 +66,7 @@
     convert_state_dict_to_kohya,
     is_wandb_available,
 )
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from diffusers.utils.import_utils import is_xformers_available
 
 
@@ -135,22 +136,7 @@ def save_model_card(
 to trigger concept `{key}` → use `{tokens}` in your prompt \n
 """
 
-    yaml = f"""---
-tags:
-- stable-diffusion
-- stable-diffusion-diffusers
-- text-to-image
-- diffusers
-- lora
-- template:sd-lora
-{img_str}
-base_model: {base_model}
-instance_prompt: {instance_prompt}
-license: openrail++
----
-"""
-
-    model_card = f"""
+    model_description = f"""
 # SD1.5 LoRA DreamBooth - {repo_id}
 
 <Gallery />
@@ -198,8 +184,29 @@ def save_model_card(
 Special VAE used for training: {vae_path}.
 
 """
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
+
+    model_card = load_or_create_model_card(
+        repo_id_or_path=repo_id,
+        from_training=True,
+        license="openrail++",
+        base_model=base_model,
+        prompt=instance_prompt,
+        model_description=model_description,
+        inference=True,
+    )
+
+    tags = [
+        "stable-diffusion",
+        "stable-diffusion-diffusers",
+        "text-to-image",
+        "diffusers",
+        "lora",
+        "template:sd-lora"
+    ]
+
+    model_card = populate_model_card(model_card, tags=tags)
+
+    model_card.save(os.path.join(repo_folder, "README.md"))
 
 
 def import_model_class_from_model_name_or_path(

examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py

@@ -69,6 +69,7 @@
     convert_unet_state_dict_to_peft,
     is_wandb_available,
 )
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.torch_utils import is_compiled_module
 
@@ -140,22 +141,7 @@ def save_model_card(
 to trigger concept `{key}` → use `{tokens}` in your prompt \n
 """
 
-    yaml = f"""---
-tags:
-- stable-diffusion-xl
-- stable-diffusion-xl-diffusers
-- text-to-image
-- diffusers
-- lora
-- template:sd-lora
-{img_str}
-base_model: {base_model}
-instance_prompt: {instance_prompt}
-license: openrail++
----
-"""
-
-    model_card = f"""
+    model_description = f"""
 # SDXL LoRA DreamBooth - {repo_id}
 
 <Gallery />
@@ -203,8 +189,29 @@ def save_model_card(
 Special VAE used for training: {vae_path}.
 
 """
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
+
+    model_card = load_or_create_model_card(
+        repo_id_or_path=repo_id,
+        from_training=True,
+        license="openrail++",
+        base_model=base_model,
+        prompt=instance_prompt,
+        model_description=model_description,
+        inference=True,
+    )
+
+    tags = [
+        "stable-diffusion-xl",
+        "stable-diffusion-xl-diffusers",
+        "text-to-image",
+        "diffusers",
+        "lora",
+        "template:sd-lora"
+    # ]
+
+    model_card = populate_model_card(model_card, tags=tags)
+
+    model_card.save(os.path.join(repo_folder, "README.md"))
 
 
 def import_model_class_from_model_name_or_path(

examples/controlnet/train_controlnet.py

@@ -49,6 +49,7 @@
 )
 from diffusers.optimization import get_scheduler
 from diffusers.utils import check_min_version, is_wandb_available
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.torch_utils import is_compiled_module
 
@@ -207,27 +208,32 @@ def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=N
             image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
             img_str += f"![images_{i})](./images_{i}.png)\n"
 
-    yaml = f"""
----
-license: creativeml-openrail-m
-base_model: {base_model}
-tags:
-- stable-diffusion
-- stable-diffusion-diffusers
-- text-to-image
-- diffusers
-- controlnet
-inference: true
----
-    """
-    model_card = f"""
+    model_description = f"""
 # controlnet-{repo_id}
 
 These are controlnet weights trained on {base_model} with new type of conditioning.
 {img_str}
 """
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
+    model_card = load_or_create_model_card(
+        repo_id_or_path=repo_id,
+        from_training=True,
+        license="creativeml-openrail-m",
+        base_model=base_model,
+        model_description=model_description,
+        inference=True,
+    )
+
+    tags = [
+        "stable-diffusion",
+        "stable-diffusion-diffusers",
+        "text-to-image",
+        "diffusers",
+        "controlnet",
+    ]
+    model_card = populate_model_card(model_card, tags=tags)
+
+    model_card.save(os.path.join(repo_folder, "README.md"))
+
 
 
 def parse_args(input_args=None):

examples/controlnet/train_controlnet_flax.py

@@ -48,7 +48,7 @@
     FlaxUNet2DConditionModel,
 )
 from diffusers.utils import check_min_version, is_wandb_available, make_image_grid
-
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 
 # To prevent an error that occurs when there are abnormally large compressed data chunk in the png image
 # see more https://github.com/python-pillow/Pillow/issues/5610
@@ -145,28 +145,33 @@ def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=N
             make_image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
             img_str += f"![images_{i})](./images_{i}.png)\n"
 
-    yaml = f"""
----
-license: creativeml-openrail-m
-base_model: {base_model}
-tags:
-- stable-diffusion
-- stable-diffusion-diffusers
-- text-to-image
-- diffusers
-- controlnet
-- jax-diffusers-event
-inference: true
----
-    """
-    model_card = f"""
+    model_description = f"""
 # controlnet- {repo_id}
 
 These are controlnet weights trained on {base_model} with new type of conditioning. You can find some example images in the following. \n
 {img_str}
 """
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
+
+    model_card = load_or_create_model_card(
+        repo_id_or_path=repo_id,
+        from_training=True,
+        license="creativeml-openrail-m",
+        base_model=base_model,
+        model_description=model_description,
+        inference=True,
+    )
+
+    tags = [
+        "stable-diffusion",
+        "stable-diffusion-diffusers",
+        "text-to-image",
+        "diffusers",
+        "controlnet",
+        "jax-diffusers-event",
+    ]
+    model_card = populate_model_card(model_card, tags=tags)
+
+    model_card.save(os.path.join(repo_folder, "README.md"))
 
 
 def parse_args():

examples/controlnet/train_controlnet_sdxl.py

@@ -51,6 +51,7 @@
 )
 from diffusers.optimization import get_scheduler
 from diffusers.utils import check_min_version, is_wandb_available, make_image_grid
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.torch_utils import is_compiled_module
 
@@ -199,28 +200,32 @@ def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=N
             make_image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png"))
             img_str += f"![images_{i})](./images_{i}.png)\n"
 
-    yaml = f"""
----
-license: openrail++
-base_model: {base_model}
-tags:
-- stable-diffusion-xl
-- stable-diffusion-xl-diffusers
-- text-to-image
-- diffusers
-- controlnet
-inference: true
----
-    """
-    model_card = f"""
+    model_description = f"""
 # controlnet-{repo_id}
 
 These are controlnet weights trained on {base_model} with new type of conditioning.
 {img_str}
 """
 
-    with open(os.path.join(repo_folder, "README.md"), "w") as f:
-        f.write(yaml + model_card)
+    model_card = load_or_create_model_card(
+        repo_id_or_path=repo_id,
+        from_training=True,
+        license="openrail++",
+        base_model=base_model,
+        model_description=model_description,
+        inference=True,
+    )
+
+    tags = [
+        "stable-diffusion",
+        "stable-diffusion-diffusers",
+        "text-to-image",
+        "diffusers",
+        "controlnet",
+    ]
+    model_card = populate_model_card(model_card, tags=tags)
+
+    model_card.save(os.path.join(repo_folder, "README.md"))
 
 
 def parse_args(input_args=None):

examples/custom_diffusion/train_custom_diffusion.py

@@ -92,7 +92,7 @@ def save_model_card(repo_id: str, images=None, base_model=str, prompt=str, repo_
         from_training=True,
         license="creativeml-openrail-m",
         base_model=base_model,
-        instance_prompt=prompt,
+        prompt=prompt,
         model_description=model_description,
         inference=True,
     )

examples/dreambooth/train_dreambooth.py

@@ -97,7 +97,7 @@ def save_model_card(
         from_training=True,
         license="creativeml-openrail-m",
         base_model=base_model,
-        instance_prompt=prompt,
+        prompt=prompt,
         model_description=model_description,
         inference=True,
     )

examples/dreambooth/train_dreambooth_lora.py

@@ -99,7 +99,7 @@ def save_model_card(
         from_training=True,
         license="creativeml-openrail-m",
         base_model=base_model,
-        instance_prompt=prompt,
+        prompt=prompt,
         model_description=model_description,
         inference=True,
     )

examples/dreambooth/train_dreambooth_lora_sdxl.py

@@ -122,7 +122,7 @@ def save_model_card(
         from_training=True,
         license="openrail++",
         base_model=base_model,
-        instance_prompt=instance_prompt,
+        prompt=instance_prompt,
         model_description=model_description,
         widget=widget_dict,
     )