Use Uploader to upload models in training time

Using two different upload methods was not a good idea.
So, stop using upload method provided by train_dreambooth_lora.py
and use Uploader class in this repo.

Also, to make it easier to port updates for train_dreambooth_lora.py
from the diffusers library, reset changes.

train_dreambooth_lora.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python
-# This file is adapted from https://github.com/huggingface/diffusers/blob/a66f2baeb782e091dde4e1e6394e46f169e5ba58/examples/dreambooth/train_dreambooth_lora.py
-# The original license is as below.
-#
 # coding=utf-8
+#
+# This file is copied from https://github.com/huggingface/diffusers/blob/febaf863026bd014b7a14349336544fc109d0f57/examples/dreambooth/train_dreambooth_lora.py
+# The original license is as below:
+#
 # Copyright 2022 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -25,6 +26,7 @@ import warnings
 from pathlib import Path
 from typing import Optional
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
@@ -48,7 +50,7 @@ from diffusers.models.cross_attention import LoRACrossAttnProcessor
 from diffusers.optimization import get_scheduler
 from diffusers.utils import check_min_version, is_wandb_available
 from diffusers.utils.import_utils import is_xformers_available
-from huggingface_hub import HfFolder, Repository, create_repo, delete_repo, whoami
+from huggingface_hub import HfFolder, Repository, create_repo, whoami
 from PIL import Image
 from torchvision import transforms
 from tqdm.auto import tqdm
@@ -61,9 +63,9 @@ check_min_version("0.12.0.dev0")
 logger = get_logger(__name__)
 
 
-def save_model_card(repo_name, base_model, instance_prompt, test_prompt="", images=None, repo_folder=""):
-    img_str = f"Test prompt: {test_prompt}\n" if test_prompt else ""
-    for i, image in enumerate(images or []):
+def save_model_card(repo_name, images=None, base_model=str, prompt=str, repo_folder=None):
+    img_str = ""
+    for i, image in enumerate(images):
         image.save(os.path.join(repo_folder, f"image_{i}.png"))
         img_str += f"![img_{i}](./image_{i}.png)\n"
 
@@ -71,7 +73,6 @@ def save_model_card(repo_name, base_model, instance_prompt, test_prompt="", imag
 ---
 license: creativeml-openrail-m
 base_model: {base_model}
-instance_prompt: {instance_prompt}
 tags:
 - stable-diffusion
 - stable-diffusion-diffusers
@@ -79,11 +80,11 @@ tags:
 - diffusers
 inference: true
 ---
-"""
+    """
     model_card = f"""
 # LoRA DreamBooth - {repo_name}
 
-These 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.\n
+These are LoRA adaption weights for {repo_name}. The weights were trained on {prompt} using [DreamBooth](https://dreambooth.github.io/). You can find some example images in the following. \n
 {img_str}
 """
     with open(os.path.join(repo_folder, "README.md"), "w") as f:
@@ -364,9 +365,6 @@ def parse_args(input_args=None):
     parser.add_argument(
         "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
     )
-    parser.add_argument("--private_repo", action="store_true")
-    parser.add_argument("--delete_existing_repo", action="store_true")
-    parser.add_argument("--upload_to_lora_library", action="store_true")
 
     if input_args is not None:
         args = parser.parse_args(input_args)
@@ -610,17 +608,11 @@ def main(args):
     if accelerator.is_main_process:
         if args.push_to_hub:
             if args.hub_model_id is None:
-                organization = 'lora-library' if args.upload_to_lora_library else None
-                repo_name = get_full_repo_name(Path(args.output_dir).name, organization=organization, token=args.hub_token)
+                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
             else:
                 repo_name = args.hub_model_id
 
-            if args.delete_existing_repo:
-                try:
-                    delete_repo(repo_name, token=args.hub_token)
-                except Exception:
-                    pass
-            create_repo(repo_name, token=args.hub_token, private=args.private_repo)
+            create_repo(repo_name, exist_ok=True, token=args.hub_token)
             repo = Repository(args.output_dir, clone_from=repo_name, token=args.hub_token)
 
             with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
@@ -826,14 +818,21 @@ def main(args):
             dirs = os.listdir(args.output_dir)
             dirs = [d for d in dirs if d.startswith("checkpoint")]
             dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-            path = dirs[-1]
-        accelerator.print(f"Resuming from checkpoint {path}")
-        accelerator.load_state(os.path.join(args.output_dir, path))
-        global_step = int(path.split("-")[1])
+            path = dirs[-1] if len(dirs) > 0 else None
 
-        resume_global_step = global_step * args.gradient_accumulation_steps
-        first_epoch = resume_global_step // num_update_steps_per_epoch
-        resume_step = resume_global_step % num_update_steps_per_epoch
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            resume_global_step = global_step * args.gradient_accumulation_steps
+            first_epoch = global_step // num_update_steps_per_epoch
+            resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)
 
     # Only show the progress bar once on each machine.
     progress_bar = tqdm(range(global_step, args.max_train_steps), disable=not accelerator.is_local_main_process)
@@ -943,6 +942,9 @@ def main(args):
             images = pipeline(prompt, num_inference_steps=25, generator=generator).images
 
             for tracker in accelerator.trackers:
+                if tracker.name == "tensorboard":
+                    np_images = np.stack([np.asarray(img) for img in images])
+                    tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC")
                 if tracker.name == "wandb":
                     tracker.log(
                         {
@@ -974,11 +976,15 @@ def main(args):
         pipeline.unet.load_attn_procs(args.output_dir)
 
         # run inference
-        generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
-        prompt = args.num_validation_images * [args.validation_prompt]
-        images = pipeline(prompt, num_inference_steps=25, generator=generator).images
+        if args.validation_prompt and args.num_validation_images > 0:
+            generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
+            prompt = args.num_validation_images * [args.validation_prompt]
+            images = pipeline(prompt, num_inference_steps=25, generator=generator).images
 
         for tracker in accelerator.trackers:
+            if tracker.name == "tensorboard":
+                np_images = np.stack([np.asarray(img) for img in images])
+                tracker.writer.add_images("test", np_images, epoch, dataformats="NHWC")
             if tracker.name == "wandb":
                 tracker.log(
                     {
@@ -992,23 +998,12 @@ def main(args):
         if args.push_to_hub:
             save_model_card(
                 repo_name,
-                base_model=args.pretrained_model_name_or_path,
-                instance_prompt=args.instance_prompt,
-                test_prompt=args.validation_prompt,
                 images=images,
-                repo_folder=args.output_dir,
-            )
-            repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)
-        else:
-            repo_name = Path(args.output_dir).name
-            save_model_card(
-                repo_name,
                 base_model=args.pretrained_model_name_or_path,
-                instance_prompt=args.instance_prompt,
-                test_prompt=args.validation_prompt,
-                images=images,
+                prompt=args.instance_prompt,
                 repo_folder=args.output_dir,
             )
+            repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)
 
     accelerator.end_training()
 

trainer.py

@@ -14,6 +14,7 @@ import torch
 from huggingface_hub import HfApi
 
 from app_upload import LoRAModelUploader
+from utils import save_model_card
 
 
 def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
@@ -125,6 +126,12 @@ class Trainer:
             command_s = ' '.join(command.split())
             f.write(command_s)
         subprocess.run(shlex.split(command))
+        save_model_card(save_dir=output_dir,
+                        base_model=base_model,
+                        instance_prompt=instance_prompt,
+                        test_prompt=validation_prompt,
+                        test_image_dir='test_images')
+
         message = 'Training completed!'
         print(message)
 

utils.py

@@ -18,3 +18,41 @@ def find_exp_dirs(ignore_repo: bool = False) -> list[str]:
             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,
+    instance_prompt: str,
+    test_prompt: str = '',
+    test_image_dir: str = '',
+) -> None:
+    image_str = ''
+    if test_prompt and test_image_dir:
+        image_paths = sorted((save_dir / test_image_dir).glob('*'))
+        if image_paths:
+            image_str = f'Test prompt: {test_prompt}\n'
+            for image_path in image_paths:
+                rel_path = image_path.relative_to(save_dir)
+                image_str += f'![{image_path.stem}]({rel_path})\n'
+
+    model_card = f'''---
+license: creativeml-openrail-m
+base_model: {base_model}
+instance_prompt: {instance_prompt}
+tags:
+- stable-diffusion
+- stable-diffusion-diffusers
+- text-to-image
+- diffusers
+inference: true
+---
+# LoRA DreamBooth - {save_dir.name}
+
+These 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.
+
+{image_str}
+'''
+
+    with open(save_dir / 'README.md', 'w') as f:
+        f.write(model_card)