From 14616f41786b2282e3686d86b2c79f9ddc95f479 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Mon, 26 Sep 2022 04:18:23 -0400
Subject: [PATCH] make metadata retrieval more robust to changes in storage
 format

- args.py will now attempt to return a metadata-containing Args
  object using the following methods:

1. By looking for the 'sd-metadata' tag in the PNG info
2. By looking from the 'Dream' tag
3. As a last resort, fetch the seed from the filename and assume
   defaults for all other options.
---
 ldm/dream/args.py | 33 +++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/ldm/dream/args.py b/ldm/dream/args.py
index 62ad9ccf01..1822e5d112 100644
--- a/ldm/dream/args.py
+++ b/ldm/dream/args.py
@@ -86,6 +86,7 @@ import shlex
 import json
 import hashlib
 import os
+import re
 import copy
 import base64
 import ldm.dream.pngwriter
@@ -732,19 +733,24 @@ def metadata_dumps(opt,
 
     return metadata
 
-def metadata_from_png(png_file_path):
+def metadata_from_png(png_file_path) -> Args:
     '''
     Given the path to a PNG file created by dream.py, retrieves
-    an Args object containing the image metadata
+    an Args object containing the image metadata. Note that this
+    returns a single Args object, not multiple.
     '''
     meta = ldm.dream.pngwriter.retrieve_metadata(png_file_path)
-    opts = metadata_loads(meta)
-    return opts[0]
+    if 'sd-metadata' in meta and len(meta['sd-metadata'])>0 :
+        return metadata_loads(meta)[0]
+    else:
+        return legacy_metadata_load(meta,png_file_path)
 
-def metadata_loads(metadata):
+def metadata_loads(metadata) ->list:
     '''
     Takes the dictionary corresponding to RFC266 (https://github.com/lstein/stable-diffusion/issues/266)
-    and returns a series of opt objects for each of the images described in the dictionary.
+    and returns a series of opt objects for each of the images described in the dictionary. Note that this
+    returns a list, and not a single object. See metadata_from_png() for a more convenient function for
+    files that contain a single image.
     '''
     results = []
     try:
@@ -797,3 +803,18 @@ def sha256(path):
             sha.update(data)
     return sha.hexdigest()
 
+def legacy_metadata_load(meta,pathname) -> Args:
+    if 'Dream' in meta and len(meta['Dream']) > 0:
+        dream_prompt = meta['Dream']
+        opt = Args()
+        opt.parse_cmd(dream_prompt)
+        return opt
+    else:               # if nothing else, we can get the seed
+        match = re.search('\d+\.(\d+)',pathname)
+        if match:
+            seed = match.groups()[0]
+            opt = Args()
+            opt.seed = seed
+            return opt
+    return None
+