fix basujindal#111: embed generation metadata in picture files.

Problems:

1. The pip package python-xmp-toolkit isn't supported on Windows,
   and doesn't automatically install the required C++-based
   dependency libexempi3. See
   https://python-xmp-toolkit.readthedocs.io/en/latest/installation.html
   for manual installation steps.

2. The metadata is just a dump of argparse's parsed options, which
   could leak private information in arguments such as outdir. A
   todo is to build a dictionary from the argparse namespace, but
   exclude fields that could contain sensitive information.

this might work

environment.yaml

@@ -24,6 +24,7 @@ dependencies:
     - transformers==4.19.2
     - torchmetrics==0.6.0
     - kornia==0.6
+    - python-xmp-toolkit==2.0.1
     - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
     - -e git+https://github.com/openai/CLIP.git@main#egg=clip
     - -e .

optimizedSD/optimized_txt2img.py

@@ -15,6 +15,8 @@
 from ldm.util import instantiate_from_config
 from optimUtils import split_weighted_subprompts, logger
 from transformers import logging
+from libxmp import XMPFiles, consts, XMPMeta  # apt install exempi; pip install python-xmp-toolkit
+
 # from samplers import CompVisDenoiser
 logging.set_verbosity_error()
 
@@ -33,6 +35,22 @@ def load_model_from_config(ckpt, verbose=False):
     return sd
 
 
+def add_metadata(filename, opt):
+    if opt.skip_metadata:
+        return
+
+    metadata = str(opt)
+
+    xmpfile = XMPFiles(file_path=filename, open_forupdate=True)
+    xmp = xmpfile.get_xmp()
+    if xmp is None:
+        xmp = XMPMeta()
+    xmp.append_array_item(consts.XMP_NS_DC, 'description', metadata,
+        {'prop_array_is_ordered':True, 'prop_value_is_array':True})
+    xmpfile.put_xmp(xmp)
+    xmpfile.close_file()
+
+
 config = "optimizedSD/v1-inference.yaml"
 ckpt = "models/ldm/stable-diffusion-v1/model.ckpt"
 
@@ -167,6 +185,11 @@ def load_model_from_config(ckpt, verbose=False):
     choices=["ddim", "plms"],
     default="plms",
 )
+parser.add_argument(
+    "--skip_metadata",
+    action='store_true',
+    help="do not add generation metadata to image file.",
+)
 opt = parser.parse_args()
 
 tic = time.time()
@@ -309,9 +332,9 @@ 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}")
-                    )
+                    filename = os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}")
+                    Image.fromarray(x_sample.astype(np.uint8)).save(filename)
+                    add_metadata(filename, opt)
                     seeds += str(opt.seed) + ","
                     opt.seed += 1
                     base_count += 1

scripts/txt2img.py

@@ -251,8 +251,7 @@ def main():
                         if not opt.skip_save:
                             for x_sample in x_samples_ddim:
                                 x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                                Image.fromarray(x_sample.astype(np.uint8)).save(
-                                    os.path.join(sample_path, f"{base_count:05}.png"))
+                                Image.fromarray(x_sample.astype(np.uint8)).save(os.path.join(sample_path, f"{base_count:05}.png"))
                                 base_count += 1
 
                         if not opt.skip_grid: