EcoCy/LoRA-DreamBooth-Training-UI

1from __future__ import annotations2 3import pathlib4 5 6def find_exp_dirs(ignore_repo: bool = False) -> list[str]:7    repo_dir = pathlib.Path(__file__).parent8    exp_root_dir = repo_dir / 'experiments'9    if not exp_root_dir.exists():10        return []11    exp_dirs = sorted(exp_root_dir.glob('*'))12    exp_dirs = [13        exp_dir for exp_dir in exp_dirs14        if (exp_dir / 'pytorch_lora_weights.bin').exists()15    ]16    if ignore_repo:17        exp_dirs = [18            exp_dir for exp_dir in exp_dirs if not (exp_dir / '.git').exists()19        ]20    return [path.relative_to(repo_dir).as_posix() for path in exp_dirs]21 22 23def save_model_card(24    save_dir: pathlib.Path,25    base_model: str,26    instance_prompt: str,27    test_prompt: str = '',28    test_image_dir: str = '',29) -> None:30    image_str = ''31    if test_prompt and test_image_dir:32        image_paths = sorted((save_dir / test_image_dir).glob('*'))33        if image_paths:34            image_str = f'Test prompt: {test_prompt}\n'35            for image_path in image_paths:36                rel_path = image_path.relative_to(save_dir)37                image_str += f'![{image_path.stem}]({rel_path})\n'38 39    model_card = f'''---40license: creativeml-openrail-m41base_model: {base_model}42instance_prompt: {instance_prompt}43tags:44- stable-diffusion45- stable-diffusion-diffusers46- text-to-image47- diffusers48- lora49inference: true50---51# LoRA DreamBooth - {save_dir.name}52 53These are LoRA adaption weights for [{base_model}](https://huggingface.co/{base_model}). The weights were trained on the instance prompt "{instance_prompt}" using [DreamBooth](https://dreambooth.github.io/). You can find some example images in the following.54 55{image_str}56'''57 58    with open(save_dir / 'README.md', 'w') as f:59        f.write(model_card)60