Save test images with train_dreambooth_lora.py

train_dreambooth_lora.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # coding=utf-8
 #
-# This file is copied from https://github.com/huggingface/diffusers/blob/febaf863026bd014b7a14349336544fc109d0f57/examples/dreambooth/train_dreambooth_lora.py
+# This file is adapted 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.
@@ -981,6 +981,12 @@ def main(args):
             prompt = args.num_validation_images * [args.validation_prompt]
             images = pipeline(prompt, num_inference_steps=25, generator=generator).images
 
+            test_image_dir = Path(args.output_dir) / 'test_images'
+            test_image_dir.mkdir()
+            for i, image in enumerate(images):
+                out_path = test_image_dir / f'image_{i}.png'
+                image.save(out_path)
+
         for tracker in accelerator.trackers:
             if tracker.name == "tensorboard":
                 np_images = np.stack([np.asarray(img) for img in images])