From 26deedc48ae76781efa1b81931c96cbbb2c154b6 Mon Sep 17 00:00:00 2001 From: Zhang Chen Date: Tue, 27 Sep 2022 00:07:37 +0800 Subject: [PATCH] Saving prompts into png metadata --- text2image.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/text2image.py b/text2image.py index 0e79fcc..69abc6b 100644 --- a/text2image.py +++ b/text2image.py @@ -2,6 +2,7 @@ from stable_diffusion_tf.stable_diffusion import Text2Image import argparse from PIL import Image +from PIL.PngImagePlugin import PngInfo parser = argparse.ArgumentParser() @@ -74,5 +75,7 @@ batch_size=1, seed=args.seed, ) -Image.fromarray(img[0]).save(args.output) +pnginfo = PngInfo() +pnginfo.add_text('prompt', args.prompt) +Image.fromarray(img[0]).save(args.output, pnginfo=pnginfo) print(f"saved at {args.output}")