From 9bef643bf50e237538cda09947c8cf299ff22612 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Fri, 21 Oct 2022 17:45:15 -0400
Subject: [PATCH] fix a few more metadata bugs

- facetool and upscale arguments now written into metadata
- cleaned up handling of !fetch command
---
 ldm/invoke/args.py      |  6 +++---
 ldm/invoke/pngwriter.py |  2 +-
 scripts/invoke.py       | 35 +++++++++++++++++------------------
 3 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/ldm/invoke/args.py b/ldm/invoke/args.py
index 2a2ee24c48..cb05ae2e58 100644
--- a/ldm/invoke/args.py
+++ b/ldm/invoke/args.py
@@ -113,8 +113,8 @@ PRECISION_CHOICES = [
 ]
 
 # is there a way to pick this up during git commits?
-APP_ID      = 'lstein/stable-diffusion'
-APP_VERSION = 'v1.15'
+APP_ID      = 'invoke-ai/InvokeAI'
+APP_VERSION = 'v2.02'
 
 class ArgFormatter(argparse.RawTextHelpFormatter):
         # use defined argument order to display usage
@@ -852,7 +852,7 @@ def metadata_dumps(opt,
     # remove any image keys not mentioned in RFC #266
     rfc266_img_fields = ['type','postprocessing','sampler','prompt','seed','variations','steps',
                          'cfg_scale','threshold','perlin','step_number','width','height','extra','strength',
-                         'init_img','init_mask']
+                         'init_img','init_mask','facetool','facetool_strength','upscale']
 
     rfc_dict ={}
 
diff --git a/ldm/invoke/pngwriter.py b/ldm/invoke/pngwriter.py
index 8f38e1add0..3a43b9bc92 100644
--- a/ldm/invoke/pngwriter.py
+++ b/ldm/invoke/pngwriter.py
@@ -38,7 +38,7 @@ class PngWriter:
         info = PngImagePlugin.PngInfo()
         info.add_text('Dream', dream_prompt)
         if metadata:
-          info.add_text('sd-metadata', json.dumps(metadata))
+            info.add_text('sd-metadata', json.dumps(metadata))
         image.save(path, 'PNG', pnginfo=info, compress_level=compress_level)
         return path
 
diff --git a/scripts/invoke.py b/scripts/invoke.py
index b2f4e35518..6acbbe63e1 100644
--- a/scripts/invoke.py
+++ b/scripts/invoke.py
@@ -809,31 +809,29 @@ def retrieve_dream_command(opt,command,completer):
         return
 
     tokens = command.split()
-    if len(tokens) > 1:
-        return write_commands(opt,tokens)
+    dir,basename = os.path.split(tokens[0])
+    if len(dir) == 0:
+        path = os.path.join(opt.outdir,basename)
+    else:
+        path = tokens[0]
 
+    if len(tokens) > 1:
+        return write_commands(opt, path, tokens[1])
+
+    cmd = ''
     try:
-        cmd = dream_cmd_from_png(tokens[0])
+        cmd = dream_cmd_from_png(path)
     except OSError:
-        print(f'## {path}: file could not be read')
+        print(f'## {tokens[0]}: file could not be read')
     except (KeyError, AttributeError, IndexError):
-        print(f'## {path}: file has no metadata')
+        print(f'## {tokens[0]}: file has no metadata')
     except:
-        print(f'## {path}: file could not be processed')
+        print(f'## {tokens[0]}: file could not be processed')
     if len(cmd)>0:
         completer.set_line(cmd)
 
-def write_commands(opt, tokens:list):
-    file_path = tokens[0]    
-    outfilepath = tokens[1]
-        
+def write_commands(opt, file_path:str, outfilepath:str):
     dir,basename = os.path.split(file_path)
-    if len(dir) == 0:
-        dir = opt.outdir
-        
-    outdir,outname = os.path.split(outfilepath)    
-    if len(outdir) == 0:
-        outfilepath = os.path.join(dir,outname)
     try:
         paths = list(Path(dir).glob(basename))
     except ValueError:
@@ -845,8 +843,6 @@ def write_commands(opt, tokens:list):
     for path in paths:
         try:
             cmd = dream_cmd_from_png(path)
-        except OSError:
-            print(f'## {path}: file could not be read')
         except (KeyError, AttributeError, IndexError):
             print(f'## {path}: file has no metadata')
         except:
@@ -855,6 +851,9 @@ def write_commands(opt, tokens:list):
             commands.append(f'# {path}')
             commands.append(cmd)
     if len(commands)>0:
+        dir,basename = os.path.split(outfilepath)
+        if len(dir)==0:
+            outfilepath = os.path.join(opt.outdir,basename)
         with open(outfilepath, 'w', encoding='utf-8') as f:
             f.write('\n'.join(commands))
         print(f'>> File {outfilepath} with commands created')