Skip to content

Commit

Permalink
[SD3-5 dreambooth lora] update model cards (#9749)
Browse files Browse the repository at this point in the history
* improve readme

* style

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
  • Loading branch information
linoytsaban and sayakpaul authored Oct 23, 2024
1 parent ab1b7b2 commit bfa0aa4
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 9 deletions.
19 changes: 14 additions & 5 deletions examples/dreambooth/train_dreambooth_lora_sd3.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,15 @@ def save_model_card(
validation_prompt=None,
repo_folder=None,
):
if "large" in base_model:
model_variant = "SD3.5-Large"
license_url = "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md"
variant_tags = ["sd3.5-large", "sd3.5", "sd3.5-diffusers"]
else:
model_variant = "SD3"
license_url = "https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE.md"
variant_tags = ["sd3", "sd3-diffusers"]

widget_dict = []
if images is not None:
for i, image in enumerate(images):
Expand All @@ -95,7 +104,7 @@ def save_model_card(
)

model_description = f"""
# SD3 DreamBooth LoRA - {repo_id}
# {model_variant} DreamBooth LoRA - {repo_id}
<Gallery />
Expand All @@ -120,7 +129,7 @@ def save_model_card(
```py
from diffusers import AutoPipelineForText2Image
import torch
pipeline = AutoPipelineForText2Image.from_pretrained('stabilityai/stable-diffusion-3-medium-diffusers', torch_dtype=torch.float16).to('cuda')
pipeline = AutoPipelineForText2Image.from_pretrained({base_model}, torch_dtype=torch.float16).to('cuda')
pipeline.load_lora_weights('{repo_id}', weight_name='pytorch_lora_weights.safetensors')
image = pipeline('{validation_prompt if validation_prompt else instance_prompt}').images[0]
```
Expand All @@ -135,7 +144,7 @@ def save_model_card(
## License
Please adhere to the licensing terms as described [here](https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE).
Please adhere to the licensing terms as described [here]({license_url}).
"""
model_card = load_or_create_model_card(
repo_id_or_path=repo_id,
Expand All @@ -151,11 +160,11 @@ def save_model_card(
"diffusers-training",
"diffusers",
"lora",
"sd3",
"sd3-diffusers",
"template:sd-lora",
]

tags += variant_tags

model_card = populate_model_card(model_card, tags=tags)
model_card.save(os.path.join(repo_folder, "README.md"))

Expand Down
16 changes: 12 additions & 4 deletions examples/dreambooth/train_dreambooth_sd3.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,15 @@ def save_model_card(
validation_prompt=None,
repo_folder=None,
):
if "large" in base_model:
model_variant = "SD3.5-Large"
license_url = "https://huggingface.co/stabilityai/stable-diffusion-3.5-large/blob/main/LICENSE.md"
variant_tags = ["sd3.5-large", "sd3.5", "sd3.5-diffusers"]
else:
model_variant = "SD3"
license_url = "https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE.md"
variant_tags = ["sd3", "sd3-diffusers"]

widget_dict = []
if images is not None:
for i, image in enumerate(images):
Expand All @@ -86,7 +95,7 @@ def save_model_card(
)

model_description = f"""
# SD3 DreamBooth - {repo_id}
# {model_variant} DreamBooth - {repo_id}
<Gallery />
Expand All @@ -113,7 +122,7 @@ def save_model_card(
## License
Please adhere to the licensing terms as described `[here](https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE)`.
Please adhere to the licensing terms as described `[here]({license_url})`.
"""
model_card = load_or_create_model_card(
repo_id_or_path=repo_id,
Expand All @@ -128,10 +137,9 @@ def save_model_card(
"text-to-image",
"diffusers-training",
"diffusers",
"sd3",
"sd3-diffusers",
"template:sd-lora",
]
tags += variant_tags

model_card = populate_model_card(model_card, tags=tags)
model_card.save(os.path.join(repo_folder, "README.md"))
Expand Down

0 comments on commit bfa0aa4

Please sign in to comment.