Skip to content

Commit

Permalink
Standardize model card for Controlnet (huggingface#6910)
Browse files Browse the repository at this point in the history
* controlnet

* style

---------

Co-authored-by: Sayak Paul <[email protected]>
  • Loading branch information
bamps53 and sayakpaul authored Feb 9, 2024
1 parent 159885a commit 59c307f
Showing 1 changed file with 21 additions and 16 deletions.
37 changes: 21 additions & 16 deletions examples/controlnet/train_controlnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -207,27 +208,31 @@ 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):
Expand Down

0 comments on commit 59c307f

Please sign in to comment.