From d33e1bf5636c55f5aea43f02ec323f16cb57a488 Mon Sep 17 00:00:00 2001
From: Sean McLellan <sean@baristalabs.io>
Date: Wed, 24 Aug 2022 12:02:36 -0400
Subject: [PATCH] Add simple way to make variants

---
 ldm/simplet2i.py |  5 +++--
 scripts/dream.py | 25 ++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py
index c7f6263816..8ffeec8dcb 100644
--- a/ldm/simplet2i.py
+++ b/ldm/simplet2i.py
@@ -143,7 +143,8 @@ The vast majority of these arguments default to reasonable values.
 
     def txt2img(self,prompt,outdir=None,batch_size=None,iterations=None,
                 steps=None,seed=None,grid=None,individual=None,width=None,height=None,
-                cfg_scale=None,ddim_eta=None,strength=None,init_img=None,skip_normalize=False):
+                cfg_scale=None,ddim_eta=None,strength=None,init_img=None,
+                skip_normalize=False,variants=None):
         """
         Generate an image from the prompt, writing iteration images into the outdir
         The output is a list of lists in the format: [[filename1,seed1], [filename2,seed2],...]
@@ -268,7 +269,7 @@ The vast majority of these arguments default to reasonable values.
     # There is lots of shared code between this and txt2img and should be refactored.
     def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=None,
                 steps=None,seed=None,grid=None,individual=None,width=None,height=None,
-                cfg_scale=None,ddim_eta=None,strength=None,skip_normalize=False):
+                cfg_scale=None,ddim_eta=None,strength=None,skip_normalize=False,variants=None):
         """
         Generate an image from the prompt and the initial image, writing iteration images into the outdir
         The output is a list of lists in the format: [[filename1,seed1], [filename2,seed2],...]
diff --git a/scripts/dream.py b/scripts/dream.py
index fb8fec2384..17b9b6c389 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -4,6 +4,7 @@ import shlex
 import atexit
 import os
 import sys
+import copy
 from PIL import Image,PngImagePlugin
 
 # readline unavailable on windows systems
@@ -162,8 +163,29 @@ def main_loop(t2i,parser,log,infile):
             results = t2i.txt2img(**vars(opt))
         else:
             results = t2i.img2img(**vars(opt))
+
+        allVariantResults = []
+        if opt.variants is not None:
+            print(f"Generating {opt.variants} variant(s)...")
+            newopt = copy.deepcopy(opt)
+            newopt.variants = None
+            for r in results:
+                resultPath = r[0]
+                print(f"\t generating variant for {resultPath}")
+                for j in range(0, opt.variants):
+                    newopt.init_img = resultPath
+                    print(f"{newopt.init_img}")
+                    variantResults = t2i.img2img(**vars(newopt))
+                    allVariantResults.append([newopt,variantResults])
+            print(f"{opt.variants} Variants generated!")
+
         print("Outputs:")
         write_log_message(t2i,opt,results,log)
+
+        if len(allVariantResults)>0:
+            print("Variant outputs:")
+            for vr in allVariantResults:
+                write_log_message(t2i,vr[0],vr[1],log)
             
 
     print("goodbye!")
@@ -285,6 +307,7 @@ def create_cmd_parser():
     parser.add_argument('-i','--individual',action='store_true',help="generate individual files (default)")
     parser.add_argument('-I','--init_img',type=str,help="path to input image (supersedes width and height)")
     parser.add_argument('-f','--strength',default=0.75,type=float,help="strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely")
+    parser.add_argument('-v','--variants',type=int,help="number of variants to generate of each image")
     parser.add_argument('-x','--skip_normalize',action='store_true',help="skip subprompt weight normalization")
     return parser
 
@@ -293,7 +316,7 @@ if readline_available:
         readline.set_completer(Completer(['cd','pwd',
                                           '--steps','-s','--seed','-S','--iterations','-n','--batch_size','-b',
                                           '--width','-W','--height','-H','--cfg_scale','-C','--grid','-g',
-                                          '--individual','-i','--init_img','-I','--strength','-f']).complete)
+                                          '--individual','-i','--init_img','-I','--strength','-f','-v','--variants']).complete)
         readline.set_completer_delims(" ")
         readline.parse_and_bind('tab: complete')
         load_history()