From fa8c67ec41ddcb895822664e31b036420b4ed58a Mon Sep 17 00:00:00 2001
From: Christian Heller <c.heller@plomlompom.de>
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