From fa8c67ec41ddcb895822664e31b036420b4ed58a Mon Sep 17 00:00:00 2001 From: Christian Heller Date: Fri, 16 Aug 2024 06:25:43 +0200 Subject: [PATCH] Initial commit. --- requirements.txt | 4 ++++ test.py | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 requirements.txt create mode 100755 test.py diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..130a1dc --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +pyexiftool +pillow +torch +diffusers diff --git a/test.py b/test.py new file mode 100755 index 0000000..72a3ebe --- /dev/null +++ b/test.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +from argparse import ArgumentParser +from random import randint +from PIL import Image +from exiftool import ExifToolHelper +from torch import Generator +from diffusers import StableDiffusionPipeline + +DEFAULT_MODEL='./v1-5-pruned-emaonly.safetensors' + +def parse_args(): + parser = ArgumentParser(add_help=False) + parser.add_argument('-p', '--prompt', required=True) + parser.add_argument('-o', '--output', required=True) + parser.add_argument('-r', '--randomness_seed', default=1, type=int, help='default 1; if set 0, chosen randomnly') + parser.add_argument('-g', '--guidance', default=7.5, type=float, help='default 7.5') + parser.add_argument('-s', '--steps', default=15, type=int, help='default 15') + 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('-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) +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']) -- 2.30.2