Allow generation of multiple image files in one go.

diff --git a/test.py b/test.py
index dff6159e82b52df4c5cee499b765a96272eada98..aeeb4766a5c457bd5819cf3f1d916b826e360600 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -5,6 +5,7 @@ from PIL import Image
 from exiftool import ExifToolHelper
 from torch import Generator
 from diffusers import StableDiffusionPipeline
+from os.path import dirname, basename, splitext, join as path_join
 
 DEFAULT_MODEL='./v1-5-pruned-emaonly.safetensors'
 
@@ -18,22 +19,29 @@ def parse_args():
     parser.add_argument('-m', '--model', default=DEFAULT_MODEL, type=str, help=f'default: {DEFAULT_MODEL}')
     parser.add_argument('-h', '--height', default=512, type=int, help='default: 512')
     parser.add_argument('-w', '--width', default=512, type=int, help='default=512')
+    parser.add_argument('-n', '--number', default=1, type=int, help='default=1')
     parser.add_argument('-H', '--help', action='help')
     return parser.parse_args()
 
 args = parse_args()
 
-generator = Generator()
-if args.randomness_seed == 0:
-    seed = randint(-(2**31-1), 2**31)
-else:
-    seed = args.randomness_seed 
-generator.manual_seed(seed)
+dir_path = dirname(args.output)
+filename = basename(args.output)
+filename_sans_ext, ext = splitext(filename)
+ext = ext[1:] if ext else 'png'
+
 pipe = StableDiffusionPipeline.from_single_file(args.model)
 pipe.to('cuda')
-image = pipe(args.prompt, generator=generator, guidance_scale=args.guidance, num_inference_steps=args.steps, height=args.height, width=args.width).images[0]
-image.save(args.output)
-metadata = f'seed: {seed}; guidance: {args.guidance}; height: {args.height}; width: {args.width}; model: {args.model}; prompt: {args.prompt}'
-print(f'saved {args.output} – metadata: {metadata}')
-with ExifToolHelper() as et:
-    et.set_tags([args.output], tags={'Comment': metadata}, params=['-overwrite_original'])
+generator = Generator()
+start_seed = randint(-(2**31-1), 2**31) if args.randomness_seed == 0 else args.randomness_seed
+for n in range(args.number):
+    seed = start_seed + n
+    generator.manual_seed(seed)
+    image = pipe(args.prompt, generator=generator, guidance_scale=args.guidance, num_inference_steps=args.steps, height=args.height, width=args.width).images[0]
+    filename_count = f'_{n:08}' if args.number > 1 else ''
+    filename = path_join(dir_path, f'{filename_sans_ext}{filename_count}.{ext}')
+    image.save(filename)
+    metadata = f'seed: {seed}; guidance: {args.guidance}; height: {args.height}; width: {args.width}; model: {args.model}; prompt: {args.prompt}'
+    print(f'saved {filename} – metadata: {metadata}')
+    with ExifToolHelper() as et:
+        et.set_tags([filename], tags={'Comment': metadata}, params=['-overwrite_original'])