|
49 | 49 | ) |
50 | 50 | from diffusers.optimization import get_scheduler |
51 | 51 | from diffusers.utils import check_min_version, is_wandb_available |
| 52 | +from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card |
52 | 53 | from diffusers.utils.import_utils import is_xformers_available |
53 | 54 | from diffusers.utils.torch_utils import is_compiled_module |
54 | 55 |
|
@@ -207,27 +208,31 @@ def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=N |
207 | 208 | image_grid(images, 1, len(images)).save(os.path.join(repo_folder, f"images_{i}.png")) |
208 | 209 | img_str += f"\n" |
209 | 210 |
|
210 | | - yaml = f""" |
211 | | ---- |
212 | | -license: creativeml-openrail-m |
213 | | -base_model: {base_model} |
214 | | -tags: |
215 | | -- stable-diffusion |
216 | | -- stable-diffusion-diffusers |
217 | | -- text-to-image |
218 | | -- diffusers |
219 | | -- controlnet |
220 | | -inference: true |
221 | | ---- |
222 | | - """ |
223 | | - model_card = f""" |
| 211 | + model_description = f""" |
224 | 212 | # controlnet-{repo_id} |
225 | 213 |
|
226 | 214 | These are controlnet weights trained on {base_model} with new type of conditioning. |
227 | 215 | {img_str} |
228 | 216 | """ |
229 | | - with open(os.path.join(repo_folder, "README.md"), "w") as f: |
230 | | - f.write(yaml + model_card) |
| 217 | + model_card = load_or_create_model_card( |
| 218 | + repo_id_or_path=repo_id, |
| 219 | + from_training=True, |
| 220 | + license="creativeml-openrail-m", |
| 221 | + base_model=base_model, |
| 222 | + model_description=model_description, |
| 223 | + inference=True, |
| 224 | + ) |
| 225 | + |
| 226 | + tags = [ |
| 227 | + "stable-diffusion", |
| 228 | + "stable-diffusion-diffusers", |
| 229 | + "text-to-image", |
| 230 | + "diffusers", |
| 231 | + "controlnet", |
| 232 | + ] |
| 233 | + model_card = populate_model_card(model_card, tags=tags) |
| 234 | + |
| 235 | + model_card.save(os.path.join(repo_folder, "README.md")) |
231 | 236 |
|
232 | 237 |
|
233 | 238 | def parse_args(input_args=None): |
|
0 commit comments