|
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