From fe12c6c099ddbc1939a06d37b9b63ed58566ef15 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Fri, 16 Sep 2022 16:58:16 -0400
Subject: [PATCH] Squashed commit of the following:

commit 67fbaa7c3106ee1c6461048cbff9111df5af0abc
Author: Lincoln Stein <lincoln.stein@gmail.com>
Date:   Fri Sep 16 16:57:54 2022 -0400

    reconciled conflicting changes to pngwriter call

commit ddc68b01f7a50901ef8d7ceb250ce4a337762819
Merge: f9feaac cbac95b
Author: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date:   Sat Sep 17 06:39:22 2022 +1000

    Merge remote-tracking branch 'upstream/development' into development

commit f9feaac8c728162f43f917364b96203fee8071c9
Author: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date:   Sat Sep 17 06:16:16 2022 +1000

    Fixes metadata related to new args

commit d1de1e357a889a350d682b03842ffca633bbf359
Author: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date:   Sat Sep 17 06:15:55 2022 +1000

    Fixes PromptFormatter import bug
---
 backend/server.py      | 20 +++++++++-----------
 ldm/dream/pngwriter.py |  9 ++++++---
 scripts/sd-metadata.py |  2 +-
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/backend/server.py b/backend/server.py
index 8dbbc65cb1..ef93c5b0d7 100644
--- a/backend/server.py
+++ b/backend/server.py
@@ -19,7 +19,8 @@ from uuid import uuid4
 from ldm.gfpgan.gfpgan_tools import real_esrgan_upscale
 from ldm.gfpgan.gfpgan_tools import run_gfpgan
 from ldm.generate import Generate
-from ldm.dream.pngwriter import PngWriter
+from ldm.dream.pngwriter import PngWriter, retrieve_metadata
+
 from modules.parameters import parameters_to_command, create_cmd_parser
 
 
@@ -146,15 +147,12 @@ def handle_request_all_images():
     paths.sort(key=lambda x: os.path.getmtime(x))
     image_array = []
     for path in paths:
-        image = Image.open(path)
-        metadata = {}
-        if 'Dream' in image.info:
-            try:
-                metadata = vars(parser.parse_args(shlex.split(image.info['Dream'])))
-            except SystemExit:
-                # TODO: Unable to parse metadata, ignore it for now,
-                # this can happen when metadata is missing a prompt
-                pass
+        # image = Image.open(path)
+        all_metadata = retrieve_metadata(path)
+        if 'Dream' in all_metadata and not all_metadata['sd-metadata']:
+            metadata = vars(parser.parse_args(shlex.split(all_metadata['Dream'])))
+        else:
+            metadata = all_metadata['sd-metadata']
         image_array.append({'path': path, 'metadata': metadata})
     return make_response("OK", data=image_array)
 
@@ -307,7 +305,7 @@ def save_image(image, parameters, output_dir, step_index=None, postprocessing=Fa
 
     command = parameters_to_command(parameters)
 
-    path = pngwriter.save_image_and_prompt_to_png(image, command, filename)
+    path = pngwriter.save_image_and_prompt_to_png(image, command, metadata=parameters, name=filename)
 
     return path
 
diff --git a/ldm/dream/pngwriter.py b/ldm/dream/pngwriter.py
index 6291059585..9a2a8bc816 100644
--- a/ldm/dream/pngwriter.py
+++ b/ldm/dream/pngwriter.py
@@ -34,6 +34,7 @@ class PngWriter:
     # saves image named _image_ to outdir/name, writing metadata from prompt
     # returns full path of output
     def save_image_and_prompt_to_png(self, image, dream_prompt, name, metadata=None):
+        print(f'self.outdir={self.outdir}, name={name}')
         path = os.path.join(self.outdir, name)
         info = PngImagePlugin.PngInfo()
         info.add_text('Dream', dream_prompt)
@@ -48,7 +49,8 @@ class PngWriter:
         metadata stored there, as a dict
         '''
         path = os.path.join(self.outdir,img_basename)
-        return retrieve_metadata(path)
+        all_metadata = retrieve_metadata(path)
+        return all_metadata['sd-metadata']
 
 def retrieve_metadata(img_path):
     '''
@@ -56,6 +58,7 @@ def retrieve_metadata(img_path):
     metadata stored there, as a dict
     '''
     im = Image.open(img_path)
-    md = im.text.get('sd-metadata',{})
-    return json.loads(md)
+    md = im.text.get('sd-metadata', '{}')
+    dream_prompt = im.text.get('Dream', '')
+    return {'sd-metadata': json.loads(md), 'Dream': dream_prompt}
 
diff --git a/scripts/sd-metadata.py b/scripts/sd-metadata.py
index a3438fa078..02d5002d60 100644
--- a/scripts/sd-metadata.py
+++ b/scripts/sd-metadata.py
@@ -13,7 +13,7 @@ filenames = sys.argv[1:]
 for f in filenames:
     try:
         metadata = retrieve_metadata(f)
-        print(f'{f}:\n',json.dumps(metadata, indent=4))
+        print(f'{f}:\n',json.dumps(metadata['sd-metadata'], indent=4))
     except FileNotFoundError:
         sys.stderr.write(f'{f} not found\n')
         continue