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'
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'])