From cdbecdb04d00fd9b1b568e76aa6561c6fbdc2123 Mon Sep 17 00:00:00 2001 From: Christian Heller Date: Fri, 16 Aug 2024 07:49:02 +0200 Subject: [PATCH] Allow generation of multiple image files in one go. --- test.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/test.py b/test.py index dff6159..aeeb476 100755 --- 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']) -- 2.30.2