from exiftool import ExifToolHelper
from torch import Generator
from diffusers import StableDiffusionPipeline
-from os.path import dirname, basename, splitext, join as path_join
+from os.path import dirname, basename, splitext, join as path_join, exists
DEFAULT_MODEL='./v1-5-pruned-emaonly.safetensors'
filename_sans_ext, ext = splitext(filename)
ext = ext[1:] if ext else 'png'
+def save_path(n: int) -> str:
+ filename_count = f'_{n:08}' if args.number > 1 else ''
+ return path_join(dir_path, f'{filename_sans_ext}{filename_count}.{ext}')
+
+for n in range(args.number):
+ path = save_path(n)
+ if exists(path):
+ raise Exception(f'Would overwrite file: {path}')
+
pipe = StableDiffusionPipeline.from_single_file(args.model)
pipe.to('cuda')
generator = Generator()
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)
+ path = save_path(n)
+ image.save(path)
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}')
+ print(f'saved {path} – metadata: {metadata}')
with ExifToolHelper() as et:
- et.set_tags([filename], tags={'Comment': metadata}, params=['-overwrite_original'])
+ et.set_tags([path], tags={'Comment': metadata}, params=['-overwrite_original'])