Merge pull request divamgupta#36 from zhangchn/png-prompt

Saving prompts into png metadata
wangcaihua · Nov 21, 2022 · c87b571 · c87b571
2 parents 93a189c + 26deedc
commit c87b571
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/text2image.py b/text2image.py
@@ -2,6 +2,7 @@
 from stable_diffusion_tf.stable_diffusion import StableDiffusion
 import argparse
 from PIL import Image
+from PIL.PngImagePlugin import PngInfo
 
 parser = argparse.ArgumentParser()
 
@@ -81,5 +82,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}")