Initial commit.

diff --git a/requirements.txt b/requirements.txt
new file mode 100644 (file)
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 (executable)
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'])