diff --git a/README.md b/README.md
index ea33dcfdc..24023ae60 100644
--- a/README.md
+++ b/README.md
@@ -119,6 +119,12 @@ This will launch gradio on port 7860 with txt2img. You can also use `docker comp
- Should generally be a multiple of 2x(n_samples)
+## `--skip_metadata`
+
+**Whether to embed generation metadata (prompt, seed, size, etc.) in the generated file.**
+
+- Uses the [XMP](https://en.wikipedia.org/wiki/Extensible_Metadata_Platform) Description field to embed most of the command-line parameters that were used to generate the image. Excludes potentially privacy-sensitive parameters such as `outdir`. Enabled by default.
+
Weighted Prompts
- Prompts can also be weighted to put relative emphasis on certain words.
diff --git a/environment.yaml b/environment.yaml
index 7f25da800..026cb3c0d 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -24,6 +24,7 @@ dependencies:
- transformers==4.19.2
- torchmetrics==0.6.0
- kornia==0.6
+ - -e git+https://github.com/sowbug/tinyxmp.git@master#egg=tinyxmp
- -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 .
diff --git a/optimizedSD/optimized_img2img.py b/optimizedSD/optimized_img2img.py
index 9ca304c2e..9238a8cad 100644
--- a/optimizedSD/optimized_img2img.py
+++ b/optimizedSD/optimized_img2img.py
@@ -17,6 +17,9 @@
from optimUtils import split_weighted_subprompts, logger
from transformers import logging
import pandas as pd
+import tinyxmp
+import xml.dom.minidom
+import json
logging.set_verbosity_error()
@@ -53,6 +56,38 @@ def load_img(path, h0, w0):
return 2.0 * image - 1.0
+def add_metadata(filename, opt):
+ if opt.skip_metadata:
+ return
+
+ SKIP_OPT_KEYS = ['outdir', 'init_img', 'from_file']
+ safe_opts = {}
+ for k, v in vars(opt).items():
+ if k in SKIP_OPT_KEYS:
+ continue
+ safe_opts[k] = v
+ metadata = json.dumps(safe_opts)
+
+ xmp_file = tinyxmp.Metadata.load(filename)
+ # Since we just generated this file, we know there's no meaningful XMP data in it.
+ # So we create an empty template.
+ xmp = '''
+
+
+
+
+
+
+
+'''
+
+ doc = xml.dom.minidom.parseString(xmp)
+ e = doc.getElementsByTagName("rdf:li")[0]
+ textnode = doc.createTextNode(metadata)
+ e.appendChild(textnode)
+ xmp_file.write_xmp(doc.childNodes[0].toxml().encode("utf-8"))
+
+
config = "optimizedSD/v1-inference.yaml"
ckpt = "models/ldm/stable-diffusion-v1/model.ckpt"
@@ -174,6 +209,11 @@ def load_img(path, h0, w0):
choices=["ddim"],
default="ddim",
)
+parser.add_argument(
+ "--skip_metadata",
+ action='store_true',
+ help="do not add generation metadata to image file.",
+)
opt = parser.parse_args()
tic = time.time()
@@ -332,9 +372,9 @@ 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}")
- )
+ 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
diff --git a/optimizedSD/optimized_txt2img.py b/optimizedSD/optimized_txt2img.py
index 6ead8618f..5c9fa3bd7 100644
--- a/optimizedSD/optimized_txt2img.py
+++ b/optimizedSD/optimized_txt2img.py
@@ -15,6 +15,10 @@
from ldm.util import instantiate_from_config
from optimUtils import split_weighted_subprompts, logger
from transformers import logging
+import tinyxmp
+import xml.dom.minidom
+import json
+
# from samplers import CompVisDenoiser
logging.set_verbosity_error()
@@ -33,6 +37,38 @@ def load_model_from_config(ckpt, verbose=False):
return sd
+def add_metadata(filename, opt):
+ if opt.skip_metadata:
+ return
+
+ SKIP_OPT_KEYS = ['outdir']
+ safe_opts = {}
+ for k, v in vars(opt).items():
+ if k in SKIP_OPT_KEYS:
+ continue
+ safe_opts[k] = v
+ metadata = json.dumps(safe_opts)
+
+ xmp_file = tinyxmp.Metadata.load(filename)
+ # Since we just generated this file, we know there's no meaningful XMP data in it.
+ # So we create an empty template.
+ xmp = '''
+
+
+
+
+
+
+
+'''
+
+ doc = xml.dom.minidom.parseString(xmp)
+ e = doc.getElementsByTagName("rdf:li")[0]
+ textnode = doc.createTextNode(metadata)
+ e.appendChild(textnode)
+ xmp_file.write_xmp(doc.childNodes[0].toxml().encode("utf-8"))
+
+
config = "optimizedSD/v1-inference.yaml"
ckpt = "models/ldm/stable-diffusion-v1/model.ckpt"
@@ -167,6 +203,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 +350,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