utils.py

from __future__ import annotations

import pathlib


def find_exp_dirs(ignore_repo: bool = False) -> list[str]:
    repo_dir = pathlib.Path(__file__).parent
    exp_root_dir = repo_dir / 'experiments'
    if not exp_root_dir.exists():
        return []
    exp_dirs = sorted(exp_root_dir.glob('*'))
    exp_dirs = [
        exp_dir for exp_dir in exp_dirs
        if (exp_dir / 'model_index.json').exists()
    ]
    if ignore_repo:
        exp_dirs = [
            exp_dir for exp_dir in exp_dirs if not (exp_dir / '.git').exists()
        ]
    return [path.relative_to(repo_dir).as_posix() for path in exp_dirs]


def save_model_card(
    save_dir: pathlib.Path,
    base_model: str,
    target_image: str,
    target_mask: str,
) -> None:
    model_card = f'''---
license: creativeml-openrail-m
base_model: {base_model}
target_image: {target_image}
target_mask: {target_mask}
tags:
- stable-diffusion-inpainting
- stable-diffusion-inpainting-diffusers
- text-to-image
- diffusers
- realfill
inference: true
---
# RealFill - {save_dir.name}

These are RealFill weights for [{base_model}](https://huggingface.co/{base_model}). The weights were trained using [RealFill](https://realfill.github.io/).
'''

    with open(save_dir / 'README.md', 'w') as f:
        f.write(model_card)