From fe6742457dc53dd662abb5c87dcfba4b4852832b Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 27 Jan 2024 15:51:04 +0530 Subject: [PATCH 01/10] feat: standarize model card creation for dreambooth training. --- examples/dreambooth/train_dreambooth.py | 67 ++++++++++++++++--------- src/diffusers/utils/hub_utils.py | 40 ++++++++++----- 2 files changed, 70 insertions(+), 37 deletions(-) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 532e134a6153..80b9d2b4b533 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -54,6 +54,7 @@ from diffusers.optimization import get_scheduler from diffusers.training_utils import compute_snr 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 @@ -69,33 +70,34 @@ def save_model_card( repo_id: str, - images=None, - base_model=str, + images: list = None, + base_model: str = None, train_text_encoder=False, - prompt=str, - repo_folder=None, + prompt: str = None, + repo_folder: str = None, pipeline: DiffusionPipeline = None, ): img_str = "" - for i, image in enumerate(images): - image.save(os.path.join(repo_folder, f"image_{i}.png")) - img_str += f"![img_{i}](./image_{i}.png)\n" - - yaml = f""" ---- -license: creativeml-openrail-m -base_model: {base_model} -instance_prompt: {prompt} -tags: -- {'stable-diffusion' if isinstance(pipeline, StableDiffusionPipeline) else 'if'} -- {'stable-diffusion-diffusers' if isinstance(pipeline, StableDiffusionPipeline) else 'if-diffusers'} -- text-to-image -- diffusers -- dreambooth -inference: true ---- - """ - model_card = f""" + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + img_str += f"![img_{i}](./image_{i}.png)\n" + + # yaml = f""" + # --- + # license: creativeml-openrail-m + # base_model: {base_model} + # instance_prompt: {prompt} + # tags: + # - {'stable-diffusion' if isinstance(pipeline, StableDiffusionPipeline) else 'if'} + # - {'stable-diffusion-diffusers' if isinstance(pipeline, StableDiffusionPipeline) else 'if-diffusers'} + # - text-to-image + # - diffusers + # - dreambooth + # inference: true + # --- + # """ + model_description = f""" # DreamBooth - {repo_id} This is a dreambooth model derived from {base_model}. The weights were trained on {prompt} using [DreamBooth](https://dreambooth.github.io/). @@ -104,8 +106,23 @@ def save_model_card( DreamBooth for the text encoder was enabled: {train_text_encoder}. """ - 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, + license="creativeml-openrail-m", + base_model=base_model, + instance_prompt=prompt, + model_description=model_description, + inference="true", + ) + + tags = ["text-to-image", "dreambooth"] + if isinstance(pipeline, StableDiffusionPipeline): + tags.extend(["stable-diffusion", "stable-diffusion-diffusers"]) + else: + tags.extend(["if", "if-diffusers"]) + model_card = populate_model_card(model_card, tags=tags) + + model_card.save(os.path.join(repo_folder, "README.md")) def log_validation( diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index c6a45b569218..69854b04a3dd 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -21,7 +21,7 @@ import traceback import warnings from pathlib import Path -from typing import Dict, Optional, Union +from typing import Dict, List, Optional, Union from uuid import uuid4 from huggingface_hub import ( @@ -94,14 +94,14 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: def load_or_create_model_card( - repo_id_or_path: Optional[str] = None, token: Optional[str] = None, is_pipeline: bool = False + repo_id_or_path: Optional[str] = None, token: Optional[str] = None, is_pipeline: bool = False, **kwargs ) -> ModelCard: """ Loads or creates a model card. Args: - repo_id (`str`): - The repo_id where to look for the model card. + repo_id_or_path (`str`): + The repo id (e.g., "runwayml/stable-diffusion-v1-5") or local path 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*): @@ -110,27 +110,43 @@ def load_or_create_model_card( 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`." + " Please make sure to have `jinja` installed before using `load_or_create_model_card`." " To install it, please run `pip install Jinja2`." ) 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() + except (EntryNotFoundError, RepositoryNotFoundError): + # Otherwise create a model card from template + if kwargs is not None: + model_description = kwargs.pop("model_description", None) + card_data = ModelCardData(**kwargs) if kwargs is not None else ModelCardData() + component = "pipeline" if is_pipeline else "model" + else: + model_description = None + + if model_description is None: + 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." + 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.""" +def populate_model_card(model_card: ModelCard, tags: Union[str, List[str]] = None) -> ModelCard: + """Populates the `model_card` with library name and optional tags.""" if model_card.data.library_name is None: model_card.data.library_name = "diffusers" + + if tags is not None: + if isinstance(tags, str): + tags = [tags] + if model_card.data.tags is None: + model_card.data.tags = [] + for tag in tags: + model_card.data.tags.append(tag) + return model_card From 8c84d91132ac6549575a04ccc6141bb2f7403d18 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 27 Jan 2024 15:55:42 +0530 Subject: [PATCH 02/10] correct 'inference --- examples/dreambooth/train_dreambooth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 80b9d2b4b533..0a295170eb1e 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -112,7 +112,7 @@ def save_model_card( base_model=base_model, instance_prompt=prompt, model_description=model_description, - inference="true", + inference=True, ) tags = ["text-to-image", "dreambooth"] From 9bc1fbc0747253663afc29cecd7754246f7f6c20 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 27 Jan 2024 16:00:51 +0530 Subject: [PATCH 03/10] remove comments. --- examples/dreambooth/train_dreambooth.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 0a295170eb1e..833f5d4c2216 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -83,20 +83,6 @@ def save_model_card( image.save(os.path.join(repo_folder, f"image_{i}.png")) img_str += f"![img_{i}](./image_{i}.png)\n" - # yaml = f""" - # --- - # license: creativeml-openrail-m - # base_model: {base_model} - # instance_prompt: {prompt} - # tags: - # - {'stable-diffusion' if isinstance(pipeline, StableDiffusionPipeline) else 'if'} - # - {'stable-diffusion-diffusers' if isinstance(pipeline, StableDiffusionPipeline) else 'if-diffusers'} - # - text-to-image - # - diffusers - # - dreambooth - # inference: true - # --- - # """ model_description = f""" # DreamBooth - {repo_id} From f463fe2d30a174190efe0c3e81fa74300c91cc09 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 27 Jan 2024 16:05:05 +0530 Subject: [PATCH 04/10] take component out of kwargs --- src/diffusers/utils/hub_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index 69854b04a3dd..98989e1fd9f3 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -119,10 +119,11 @@ def load_or_create_model_card( model_card = ModelCard.load(repo_id_or_path, token=token) except (EntryNotFoundError, RepositoryNotFoundError): # Otherwise create a model card from template + component = "pipeline" if is_pipeline else "model" + if kwargs is not None: model_description = kwargs.pop("model_description", None) card_data = ModelCardData(**kwargs) if kwargs is not None else ModelCardData() - component = "pipeline" if is_pipeline else "model" else: model_description = None From dc9afd38da564cf0d996f80e498946a22132a83e Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 27 Jan 2024 16:06:00 +0530 Subject: [PATCH 05/10] style --- src/diffusers/utils/hub_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index 98989e1fd9f3..9a5dafd7ec5a 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -120,7 +120,7 @@ def load_or_create_model_card( except (EntryNotFoundError, RepositoryNotFoundError): # Otherwise create a model card from template component = "pipeline" if is_pipeline else "model" - + if kwargs is not None: model_description = kwargs.pop("model_description", None) card_data = ModelCardData(**kwargs) if kwargs is not None else ModelCardData() From 52d8131cf566d62b0f87b81defb9142512f08b18 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 5 Feb 2024 10:54:09 +0530 Subject: [PATCH 06/10] add: card template to have a leaner description. --- examples/dreambooth/train_dreambooth.py | 1 + src/diffusers/utils/hub_utils.py | 50 ++++++++++++++++------ src/diffusers/utils/model_card_template.md | 24 +++++++++++ 3 files changed, 62 insertions(+), 13 deletions(-) create mode 100644 src/diffusers/utils/model_card_template.md diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 833f5d4c2216..4d899ef56e37 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -94,6 +94,7 @@ def save_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, instance_prompt=prompt, diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index 9a5dafd7ec5a..622a6b91529e 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -65,7 +65,7 @@ logger = get_logger(__name__) - +MODEL_CARD_TEMPLATE_PATH = Path(__file__).parent / "model_card_template.md" SESSION_ID = uuid4().hex @@ -94,7 +94,15 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: def load_or_create_model_card( - repo_id_or_path: Optional[str] = None, token: Optional[str] = None, is_pipeline: bool = False, **kwargs + repo_id_or_path: Optional[str] = None, + token: Optional[str] = None, + is_pipeline: bool = False, + from_training: bool = False, + model_description: str = None, + base_model: str = None, + instance_prompt: str = None, + license: str = None, + inference: bool = None, ) -> ModelCard: """ Loads or creates a model card. @@ -106,6 +114,16 @@ def load_or_create_model_card( 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`]. + from_training: (`bool`, *optional*): Boolean flag to denote if the model card is being created from a training script. + model_description (`str`, *optional*): Model description to add to the model card. Helpful when using + `load_or_create_model_card` from a training script. + base_model (`str`, *optional*): Base model identifier (e.g., "stabilityai/stable-diffusion-xl-base-1.0"). Useful + for DreamBooth-like training. + instance_prompt (`str`, *optional*): Instance prompt used for training. Useful for DreamBooth-like training. + license: (`str`, *optional*): License of the output artifact. Helpful when using + `load_or_create_model_card` from a training script. + inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using + `load_or_create_model_card` from a training script. """ if not is_jinja_available(): raise ValueError( @@ -119,18 +137,24 @@ def load_or_create_model_card( model_card = ModelCard.load(repo_id_or_path, token=token) except (EntryNotFoundError, RepositoryNotFoundError): # Otherwise create a model card from template - component = "pipeline" if is_pipeline else "model" - - if kwargs is not None: - model_description = kwargs.pop("model_description", None) - card_data = ModelCardData(**kwargs) if kwargs is not None else ModelCardData() + if from_training: + model_card = ModelCard.from_template( + card_data=ModelCardData( # Card metadata object that will be converted to YAML block + license=license, + library_name="diffusers", + inference=inference, + base_model=base_model, + instance_prompt=instance_prompt, + ), + template_path=MODEL_CARD_TEMPLATE_PATH, + model_description=model_description, + ) else: - model_description = None - - if model_description is None: - 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." - - model_card = ModelCard.from_template(card_data, model_description=model_description) + card_data = ModelCardData() + component = "pipeline" if is_pipeline else "model" + if model_description is None: + 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." + model_card = ModelCard.from_template(card_data, model_description=model_description) return model_card diff --git a/src/diffusers/utils/model_card_template.md b/src/diffusers/utils/model_card_template.md new file mode 100644 index 000000000000..f41b71e24e20 --- /dev/null +++ b/src/diffusers/utils/model_card_template.md @@ -0,0 +1,24 @@ +--- +{{ card_data }} +--- + + + +{{ model_description }} + +## Intended uses & limitations + +#### How to use + +```python +# TODO: add an example code snippet for running this diffusion pipeline +``` + +#### Limitations and bias + +[TODO: provide examples of latent issues and potential remediations] + +## Training details + +[TODO: describe the data used to train the model] From 2f647ddd7ff5070f9f85ed9d28e08caed1c3f5b9 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 5 Feb 2024 14:50:11 +0530 Subject: [PATCH 07/10] widget support. --- .../dreambooth/train_dreambooth_lora_sdxl.py | 64 ++++++++++--------- src/diffusers/utils/hub_utils.py | 3 + 2 files changed, 36 insertions(+), 31 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py index aa09bf9a0ebf..f407f88852fe 100644 --- a/examples/dreambooth/train_dreambooth_lora_sdxl.py +++ b/examples/dreambooth/train_dreambooth_lora_sdxl.py @@ -62,6 +62,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 @@ -75,40 +76,22 @@ def save_model_card( repo_id: str, images=None, - base_model=str, + base_model: str = None, train_text_encoder=False, - instance_prompt=str, - validation_prompt=str, + instance_prompt=None, + validation_prompt=None, repo_folder=None, vae_path=None, ): - img_str = "widget:\n" if images else "" - for i, image in enumerate(images): - image.save(os.path.join(repo_folder, f"image_{i}.png")) - img_str += f""" - - text: '{validation_prompt if validation_prompt else ' ' }' - output: - url: - "image_{i}.png" - """ - - 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++ ---- - """ + widget_dict = [] + if images is not None: + for i, image in enumerate(images): + image.save(os.path.join(repo_folder, f"image_{i}.png")) + widget_dict.append( + {"text": validation_prompt if validation_prompt else " ", "output": {"url": f"image_{i}.png"}} + ) - model_card = f""" + model_description = f""" # SDXL LoRA DreamBooth - {repo_id} @@ -134,8 +117,27 @@ def save_model_card( [Download]({repo_id}/tree/main) them in the Files & versions tab. """ - 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, + instance_prompt=instance_prompt, + model_description=model_description, + widget_str=widget_dict, + ) + tags = [ + "text-to-image", + "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( diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index 622a6b91529e..bddd14c3fd6a 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -102,6 +102,7 @@ def load_or_create_model_card( base_model: str = None, instance_prompt: str = None, license: str = None, + widget_str: str = None, inference: bool = None, ) -> ModelCard: """ @@ -122,6 +123,7 @@ def load_or_create_model_card( instance_prompt (`str`, *optional*): Instance prompt used for training. Useful for DreamBooth-like training. license: (`str`, *optional*): License of the output artifact. Helpful when using `load_or_create_model_card` from a training script. + widget_str (`str`, *optional*): Widget string to accompany a gallery template. inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using `load_or_create_model_card` from a training script. """ @@ -145,6 +147,7 @@ def load_or_create_model_card( inference=inference, base_model=base_model, instance_prompt=instance_prompt, + widget=widget_str, ), template_path=MODEL_CARD_TEMPLATE_PATH, model_description=model_description, From cc1c73ed071c876de9d17c51ec8b53e0b5db87d2 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 5 Feb 2024 14:56:48 +0530 Subject: [PATCH 08/10] propagate changes to train_dreambooth_lora --- examples/dreambooth/train_dreambooth_lora.py | 36 +++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 3724e3d140d9..f0c47821b0c9 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -61,6 +61,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 @@ -85,21 +86,7 @@ def save_model_card( image.save(os.path.join(repo_folder, f"image_{i}.png")) img_str += f"![img_{i}](./image_{i}.png)\n" - yaml = f""" ---- -license: creativeml-openrail-m -base_model: {base_model} -instance_prompt: {prompt} -tags: -- {'stable-diffusion' if isinstance(pipeline, StableDiffusionPipeline) else 'if'} -- {'stable-diffusion-diffusers' if isinstance(pipeline, StableDiffusionPipeline) else 'if-diffusers'} -- text-to-image -- diffusers -- lora -inference: true ---- - """ - model_card = f""" + model_description = f""" # LoRA DreamBooth - {repo_id} These are LoRA adaption weights for {base_model}. The weights were trained on {prompt} using [DreamBooth](https://dreambooth.github.io/). You can find some example images in the following. \n @@ -107,8 +94,23 @@ def save_model_card( LoRA for the text encoder was enabled: {train_text_encoder}. """ - 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, + instance_prompt=prompt, + model_description=model_description, + inference=True, + ) + tags = ["text-to-image", "diffusers", "lora"] + if isinstance(pipeline, StableDiffusionPipeline): + tags.extend(["stable-diffusion", "stable-diffusion-diffusers"]) + else: + tags.extend(["if", "if-diffusers"]) + 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(pretrained_model_name_or_path: str, revision: str): From 2650fed7953d2fc5a85c0fb5c88e184e0861a059 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 5 Feb 2024 15:00:00 +0530 Subject: [PATCH 09/10] propagate changes to custom diffusion --- .../train_custom_diffusion.py | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/examples/custom_diffusion/train_custom_diffusion.py b/examples/custom_diffusion/train_custom_diffusion.py index 559430eba177..0ea2732a959f 100644 --- a/examples/custom_diffusion/train_custom_diffusion.py +++ b/examples/custom_diffusion/train_custom_diffusion.py @@ -58,6 +58,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 @@ -78,21 +79,7 @@ def save_model_card(repo_id: str, images=None, base_model=str, prompt=str, repo_ image.save(os.path.join(repo_folder, f"image_{i}.png")) img_str += f"![img_{i}](./image_{i}.png)\n" - yaml = f""" ---- -license: creativeml-openrail-m -base_model: {base_model} -instance_prompt: {prompt} -tags: -- stable-diffusion -- stable-diffusion-diffusers -- text-to-image -- diffusers -- custom-diffusion -inference: true ---- - """ - model_card = f""" + model_description = f""" # Custom Diffusion - {repo_id} These are Custom Diffusion adaption weights for {base_model}. The weights were trained on {prompt} using [Custom Diffusion](https://www.cs.cmu.edu/~custom-diffusion). You can find some example images in the following. \n @@ -100,8 +87,20 @@ def save_model_card(repo_id: str, images=None, base_model=str, prompt=str, repo_ \nFor more details on the training, please follow [this link](https://github.com/huggingface/diffusers/blob/main/examples/custom_diffusion). """ - 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, + instance_prompt=prompt, + model_description=model_description, + inference=True, + ) + + tags = ["text-to-image", "diffusers", "stable-diffusion", "stable-diffusion-diffusers", "custom-diffusion"] + 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(pretrained_model_name_or_path: str, revision: str): From 2c0a3cf944195f50085bfbaf200c3144dc361d16 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 7 Feb 2024 12:06:07 +0530 Subject: [PATCH 10/10] make widget properly type-annotated --- .../dreambooth/train_dreambooth_lora_sdxl.py | 2 +- src/diffusers/utils/hub_utils.py | 26 +++++++++---------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py index f407f88852fe..8df61f132510 100644 --- a/examples/dreambooth/train_dreambooth_lora_sdxl.py +++ b/examples/dreambooth/train_dreambooth_lora_sdxl.py @@ -124,7 +124,7 @@ def save_model_card( base_model=base_model, instance_prompt=instance_prompt, model_description=model_description, - widget_str=widget_dict, + widget=widget_dict, ) tags = [ "text-to-image", diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index bddd14c3fd6a..7528e34d3332 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -94,16 +94,16 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str: def load_or_create_model_card( - repo_id_or_path: Optional[str] = None, + repo_id_or_path: str = None, token: Optional[str] = None, is_pipeline: bool = False, from_training: bool = False, - model_description: str = None, + model_description: Optional[str] = None, base_model: str = None, - instance_prompt: str = None, - license: str = None, - widget_str: str = None, - inference: bool = None, + prompt: Optional[str] = None, + license: Optional[str] = None, + widget: Optional[List[dict]] = None, + inference: Optional[bool] = None, ) -> ModelCard: """ Loads or creates a model card. @@ -113,17 +113,17 @@ def load_or_create_model_card( The repo id (e.g., "runwayml/stable-diffusion-v1-5") or local path 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*): + is_pipeline (`bool`): Boolean to indicate if we're adding tag to a [`DiffusionPipeline`]. - from_training: (`bool`, *optional*): Boolean flag to denote if the model card is being created from a training script. + from_training: (`bool`): Boolean flag to denote if the model card is being created from a training script. model_description (`str`, *optional*): Model description to add to the model card. Helpful when using `load_or_create_model_card` from a training script. - base_model (`str`, *optional*): Base model identifier (e.g., "stabilityai/stable-diffusion-xl-base-1.0"). Useful + base_model (`str`): Base model identifier (e.g., "stabilityai/stable-diffusion-xl-base-1.0"). Useful for DreamBooth-like training. - instance_prompt (`str`, *optional*): Instance prompt used for training. Useful for DreamBooth-like training. + prompt (`str`, *optional*): Prompt used for training. Useful for DreamBooth-like training. license: (`str`, *optional*): License of the output artifact. Helpful when using `load_or_create_model_card` from a training script. - widget_str (`str`, *optional*): Widget string to accompany a gallery template. + widget (`List[dict]`, *optional*): Widget to accompany a gallery template. inference: (`bool`, optional): Whether to turn on inference widget. Helpful when using `load_or_create_model_card` from a training script. """ @@ -146,8 +146,8 @@ def load_or_create_model_card( library_name="diffusers", inference=inference, base_model=base_model, - instance_prompt=instance_prompt, - widget=widget_str, + instance_prompt=prompt, + widget=widget, ), template_path=MODEL_CARD_TEMPLATE_PATH, model_description=model_description,