From d63897fc39d9e1aa4e8b6a7a57071265b5cecdb8 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Mon, 26 Sep 2022 03:23:03 -0400
Subject: [PATCH] tweaks to get embiggen to work with new development changes
 PR #553

- swapped out calls to choose_autocast()
- fixed call to retrieve_png_metadata() so as to recover original prompt

Original author: @travco
---
 ldm/dream/generator/embiggen.py | 24 +++++-------------------
 ldm/generate.py                 |  7 ++-----
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/ldm/dream/generator/embiggen.py b/ldm/dream/generator/embiggen.py
index eaca9a0416..c85bae868c 100644
--- a/ldm/dream/generator/embiggen.py
+++ b/ldm/dream/generator/embiggen.py
@@ -7,9 +7,9 @@ import torch
 import numpy as  np
 from tqdm import trange
 from PIL               import Image
-from ldm.dream.devices import choose_autocast_device
 from ldm.dream.generator.base      import Generator
 from ldm.dream.generator.img2img   import Img2Img
+from ldm.dream.devices import choose_autocast
 
 class Embiggen(Generator):
     def __init__(self, model, precision):
@@ -20,16 +20,17 @@ class Embiggen(Generator):
     def generate(self,prompt,iterations=1,seed=None,
                  image_callback=None, step_callback=None,
                  **kwargs):
-        device_type,scope   = choose_autocast_device(self.model.device)
-        make_image          = self.get_make_image(
+        scope      = choose_autocast(self.precision)
+        make_image = self.get_make_image(
             prompt,
             step_callback = step_callback,
             **kwargs
         )
         results             = []
         seed                = seed if seed else self.new_seed()
+
         # Noise will be generated by the Img2Img generator when called
-        with scope(device_type), self.model.ema_scope():
+        with scope(self.model.device.type), self.model.ema_scope():
             for n in trange(iterations, desc='Generating'):
                 # make_image will call Img2Img which will do the equivalent of get_noise itself
                 image = make_image()
@@ -360,21 +361,6 @@ class Embiggen(Generator):
                     init_image     = newinitimage,    # notice that init_image is different from init_img
                     mask_image     = None,
                     strength       = strength,
-                    iterations=1,
-                    seed=self.seed,
-                    sampler=sampler,
-                    steps=steps,
-                    cfg_scale=cfg_scale,
-                    conditioning=conditioning,
-                    ddim_eta=ddim_eta,
-                    image_callback=None,  # called only after the final image is generated
-                    step_callback=step_callback,   # called after each intermediate image is generated
-                    width=width,
-                    height=height,
-                    init_img=init_img,        # img2img doesn't need this, but it might in the future
-                    init_image=newinitimage,    # notice that init_image is different from init_img
-                    mask_image=None,
-                    strength=strength,
                 )
 
                 emb_tile_store.append(tile_results[0][0])
diff --git a/ldm/generate.py b/ldm/generate.py
index 854a519520..a283645225 100644
--- a/ldm/generate.py
+++ b/ldm/generate.py
@@ -497,11 +497,8 @@ class Generate:
         prompt = None
         try:
             args = metadata_from_png(image_path)
-            if len(args) > 1:
-                print("* Can't postprocess a grid")
-                return
-            seed   = args[0].seed
-            prompt = args[0].prompt
+            seed   = args.seed
+            prompt = args.prompt
             print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}')
         except:
             m    = re.search('(\d+)\.png$',image_path)