From 48aa6416dc194d0caf4ddff988d953a42dc585b1 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Tue, 8 Nov 2022 15:22:32 +0000
Subject: [PATCH] enable outcropping of random JPG/PNG images

- Works best with runwayML inpainting model
- Numerous code changes required to propagate seed to final metadata.
  Original code predicated on the image being generated within InvokeAI.
---
 ldm/generate.py                   | 13 +++++++------
 ldm/invoke/generator/base.py      |  2 +-
 ldm/invoke/prompt_parser.py       |  2 +-
 ldm/invoke/restoration/outcrop.py |  5 +++--
 scripts/invoke.py                 |  7 +++++--
 5 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/ldm/generate.py b/ldm/generate.py
index e2d4a40de7..746e56dd11 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -566,12 +566,13 @@ class Generate:
 
         args   = metadata_from_png(image_path)
         seed   = args.seed
-        prompt = args.prompt
-        print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}')
-
-        if not seed:
-            print('* Could not recover seed for image. Replacing with 42. This will not affect image quality')
-            seed = 42
+        prompt = args.prompt or ''
+        if seed == 0:
+            seed = random.randrange(0, np.iinfo(np.uint32).max)
+            opt.seed = seed
+            print(f'>> generated new seed {seed} and prompt "{prompt}" for {image_path}')
+        else:
+            print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}')
 
         # try to reuse the same filename prefix as the original file.
         # we take everything up to the first period
diff --git a/ldm/invoke/generator/base.py b/ldm/invoke/generator/base.py
index 3c6eca08a2..719d08c7c0 100644
--- a/ldm/invoke/generator/base.py
+++ b/ldm/invoke/generator/base.py
@@ -63,7 +63,7 @@ class Generator():
             **kwargs
         )
         results             = []
-        seed                = seed if seed is not None else self.new_seed()
+        seed                = seed if seed is not None and seed > 0 else self.new_seed()
         first_seed          = seed
         seed, initial_noise = self.generate_initial_noise(seed, width, height)
 
diff --git a/ldm/invoke/prompt_parser.py b/ldm/invoke/prompt_parser.py
index 3dbcc1bb4b..42c83188aa 100644
--- a/ldm/invoke/prompt_parser.py
+++ b/ldm/invoke/prompt_parser.py
@@ -636,7 +636,7 @@ def split_weighted_subprompts(text, skip_normalize=False)->list:
     weight_sum = sum(map(lambda x: x[1], parsed_prompts))
     if weight_sum == 0:
         print(
-            "Warning: Subprompt weights add up to zero. Discarding and using even weights instead.")
+            "* Warning: Subprompt weights add up to zero. Discarding and using even weights instead.")
         equal_weight = 1 / max(len(parsed_prompts), 1)
         return [(x[0], equal_weight) for x in parsed_prompts]
     return [(x[0], x[1] / weight_sum) for x in parsed_prompts]
diff --git a/ldm/invoke/restoration/outcrop.py b/ldm/invoke/restoration/outcrop.py
index b5d42250c5..1a0aaf2c8f 100644
--- a/ldm/invoke/restoration/outcrop.py
+++ b/ldm/invoke/restoration/outcrop.py
@@ -28,11 +28,12 @@ class Outcrop(object):
         self.generate._set_sampler()
 
         def wrapped_callback(img,seed,**kwargs):
-            image_callback(img,orig_opt.seed,use_prefix=prefix,**kwargs)
+            preferred_seed = orig_opt.seed if orig_opt.seed> 0 else seed
+            image_callback(img,preferred_seed,use_prefix=prefix,**kwargs)
 
         result= self.generate.prompt2image(
             orig_opt.prompt,
-            seed        = orig_opt.seed,    # uncomment to make it deterministic
+            seed        = orig_opt.seed if orig_opt.seed>0 else opt.seed,
             sampler     = self.generate.sampler,
             steps       = opt.steps,
             cfg_scale   = opt.cfg_scale,
diff --git a/scripts/invoke.py b/scripts/invoke.py
index 1e9a84295e..325a3ca6c4 100755
--- a/scripts/invoke.py
+++ b/scripts/invoke.py
@@ -29,6 +29,7 @@ infile = None
 def main():
     """Initialize command-line parsers and the diffusion model"""
     global infile
+    print('* Initializing, be patient...')
     
     opt  = Args()
     args = opt.parse_args()
@@ -46,7 +47,6 @@ def main():
             print('--max_loaded_models must be >= 1; using 1')
             args.max_loaded_models = 1
 
-    print('* Initializing, be patient...')
     from ldm.generate import Generate
 
     # these two lines prevent a horrible warning message from appearing
@@ -656,7 +656,10 @@ def do_postprocess (gen, opt, callback):
 def add_postprocessing_to_metadata(opt,original_file,new_file,tool,command):
     original_file = original_file if os.path.exists(original_file) else os.path.join(opt.outdir,original_file)
     new_file       = new_file     if os.path.exists(new_file)      else os.path.join(opt.outdir,new_file)
-    meta = retrieve_metadata(original_file)['sd-metadata']
+    try:
+        meta = retrieve_metadata(original_file)['sd-metadata']
+    except AttributeError:
+        meta = retrieve_metadata(new_file)['sd-metadata']
     if 'image' not in meta:
         meta = metadata_dumps(opt,seeds=[opt.seed])['image']
         meta['image'] = {}