Adds Exif and PNG Info to all scripts

optimizedSD/img2img_gradio.py

@@ -21,7 +21,8 @@
 from ldm.util import instantiate_from_config
 from transformers import logging
 import pandas as pd
-from optimUtils import split_weighted_subprompts, logger
+from optimUtils import split_weighted_subprompts, logger, create_exif_info
+
 logging.set_verbosity_error()
 import mimetypes
 mimetypes.init()
@@ -221,9 +222,12 @@ def generate(
                         x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                         all_samples.append(x_sample.to("cpu"))
                         x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c")
-                        Image.fromarray(x_sample.astype(np.uint8)).save(
-                            os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}")
-                        )
+                        img = Image.fromarray(x_sample.astype(np.uint8))
+                        target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" +
+                                                   f"{base_count:05}.{img_format}")
+                        info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale,
+                                                      full_precision, batch_size, i, "img2img_gradio.py")
+                        img.save(target_path, pnginfo=info, exif=exif)
                         seeds += str(seed) + ","
                         seed += 1
                         base_count += 1

optimizedSD/inpaint_gradio.py

@@ -19,7 +19,7 @@
 from transformers import logging
 
 from ldm.util import instantiate_from_config
-from optimUtils import split_weighted_subprompts, logger
+from optimUtils import split_weighted_subprompts, logger, create_exif_info
 
 logging.set_verbosity_error()
 import mimetypes
@@ -220,9 +220,12 @@ def generate(
                         x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                         all_samples.append(x_sample.to("cpu"))
                         x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c")
-                        Image.fromarray(x_sample.astype(np.uint8)).save(
-                            os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}")
-                        )
+                        img = Image.fromarray(x_sample.astype(np.uint8))
+                        target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" +
+                                                   f"{base_count:05}.{img_format}")
+                        info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale,
+                                                      full_precision, batch_size, i, "inpaint_gradio.py")
+                        img.save(target_path, pnginfo=info, exif=exif)
                         seeds += str(seed) + ","
                         seed += 1
                         base_count += 1

optimizedSD/optimUtils.py

@@ -1,5 +1,7 @@
 import os
 import pandas as pd
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
 
 
 def split_weighted_subprompts(text):
@@ -51,7 +53,7 @@ def logger(params, log_csv):
     os.makedirs('logs', exist_ok=True)
     cols = [arg for arg, _ in params.items()]
     if not os.path.exists(log_csv):
-        df = pd.DataFrame(columns=cols) 
+        df = pd.DataFrame(columns=cols)
         df.to_csv(log_csv, index=False)
 
     df = pd.read_csv(log_csv)
@@ -70,4 +72,29 @@ def logger(params, log_csv):
             li[col] = ''
 
     df = pd.DataFrame(li,index = [0])
-    df.to_csv(log_csv,index=False, mode='a', header=False)
+    df.to_csv(log_csv,index=False, mode='a', header=False)
+
+
+def create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, full_precision, batch_size,
+                     batch_index, scriptname):
+    info = PngInfo()
+    info.add_itxt("Software", f"Stable Diffusion (github.com/basujindal/stable-diffusion), "
+                              f"{scriptname}", zip=True)
+    info.add_itxt("Seed", str(seed))
+    info.add_itxt("Prompt", prompt, zip=True)
+    info.add_itxt("ddim_steps", str(ddim_steps))
+    info.add_itxt("sampler", str(sampler))
+    info.add_itxt("Scale", str(scale))
+    info.add_itxt("FullPrecision", str(full_precision))
+    info.add_itxt("Batch Size", str(batch_size))
+    info.add_itxt("Batch Index", str(batch_index))
+    exif = img.getexif()
+    exif[0xA401] = 1  # Yes, this has had special processing on it
+    exif[0xA40B] = f"Seed: {seed}, ddim_steps: {ddim_steps}, sampler: {sampler}, scale: {scale}, " \
+                   f"precision: {str(full_precision)}, batchsize: {batch_size}, batchindex: {batch_index}, " \
+                   f"prompt: {prompt}"
+    exif[0x927C] = exif[0xA40B]
+    exif[0x9286] = exif[0xA40B]
+    exif[0x010E] = prompt
+    exif[0x0131] = f"Stable Diffusion (github.com/basujindal/stable-diffusion) {scriptname}"
+    return info, exif

optimizedSD/optimized_img2img.py

@@ -14,7 +14,7 @@
 from contextlib import contextmanager, nullcontext
 from einops import rearrange, repeat
 from ldm.util import instantiate_from_config
-from optimUtils import split_weighted_subprompts, logger
+from optimUtils import split_weighted_subprompts, logger, create_exif_info
 from transformers import logging
 import pandas as pd
 logging.set_verbosity_error()
@@ -243,7 +243,8 @@ def load_img(path, h0, w0):
 else:
     print(f"reading prompts from {opt.from_file}")
     with open(opt.from_file, "r") as f:
-        data = f.read().splitlines()
+        prompt = f.read()
+        data = prompt.splitlines()
         data = batch_size * list(data)
         data = list(chunk(sorted(data), batch_size))
 
@@ -332,8 +333,12 @@ def load_img(path, h0, w0):
                     x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))
                     x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                     x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c")
-                    Image.fromarray(x_sample.astype(np.uint8)).save(
-                        os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}")
+                    img = Image.fromarray(x_sample.astype(np.uint8))
+                    info, exif = create_exif_info(img, opt.seed, prompt, opt.ddim_steps, opt.sampler, opt.scale,
+                                                  opt.precision, batch_size, i, "optimized_img2img.py")
+                    img.save(
+                        os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}"),
+                        pnginfo=info
                     )
                     seeds += str(opt.seed) + ","
                     opt.seed += 1

optimizedSD/optimized_txt2img.py

@@ -13,7 +13,7 @@
 from torch import autocast
 from contextlib import contextmanager, nullcontext
 from ldm.util import instantiate_from_config
-from optimUtils import split_weighted_subprompts, logger
+from optimUtils import split_weighted_subprompts, logger, create_exif_info
 from transformers import logging
 # from samplers import CompVisDenoiser
 logging.set_verbosity_error()
@@ -237,7 +237,8 @@ def load_model_from_config(ckpt, verbose=False):
 else:
     print(f"reading prompts from {opt.from_file}")
     with open(opt.from_file, "r") as f:
-        data = f.read().splitlines()
+        prompt = f.read()
+        data = prompt.splitlines()
         data = batch_size * list(data)
         data = list(chunk(sorted(data), batch_size))
 
@@ -309,8 +310,12 @@ def load_model_from_config(ckpt, verbose=False):
                     x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))
                     x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                     x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c")
-                    Image.fromarray(x_sample.astype(np.uint8)).save(
-                        os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}")
+                    img = Image.fromarray(x_sample.astype(np.uint8))
+                    info, exif = create_exif_info(img, opt.seed, prompt, opt.ddim_steps, opt.sampler, opt.scale,
+                                                  opt.precision, batch_size, i, "optimized_txt2img.py")
+                    img.save(
+                        os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}"),
+                        pnginfo=info
                     )
                     seeds += str(opt.seed) + ","
                     opt.seed += 1

optimizedSD/txt2img_gradio.py

@@ -20,7 +20,7 @@
 from torch import autocast
 from contextlib import nullcontext
 from ldm.util import instantiate_from_config
-from optimUtils import split_weighted_subprompts, logger
+from optimUtils import split_weighted_subprompts, logger, create_exif_info
 from transformers import logging
 logging.set_verbosity_error()
 import mimetypes
@@ -95,7 +95,6 @@ def generate(
     full_precision,
     sampler,
 ):
-
     C = 4
     f = 8
     start_code = None
@@ -178,20 +177,22 @@ def generate(
                         unconditional_conditioning=uc,
                         eta=ddim_eta,
                         x_T=start_code,
-                        sampler = sampler,
+                        sampler=sampler,
                     )
 
                     modelFS.to(device)
                     print("saving images")
                     for i in range(batch_size):
-
                         x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0))
                         x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                         all_samples.append(x_sample.to("cpu"))
                         x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c")
-                        Image.fromarray(x_sample.astype(np.uint8)).save(
-                            os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}")
-                        )
+                        img = Image.fromarray(x_sample.astype(np.uint8))
+                        target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" +
+                                                   f"{base_count:05}.{img_format}")
+                        info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale,
+                                                      full_precision, batch_size, i, "text2img_gradio.py")
+                        img.save(target_path, pnginfo=info, exif=exif)
                         seeds += str(seed) + ","
                         seed += 1
                         base_count += 1