feature complete; looks like ready for merge

TODO.txt

@@ -2,6 +2,7 @@ Feature requests:
 
 1. "gobig" mode - split image into strips, scale up, add detail using
    img2img and reassemble with feathering. Issue #66.
+   See https://github.com/jquesnelle/txt2imghd
 
 2. Port basujindal low VRAM optimizations. Issue #62
 

ldm/dream_util.py

@@ -2,6 +2,7 @@
 import os
 import atexit
 import re
+from math import sqrt,floor,ceil
 from PIL import Image,PngImagePlugin
 
 # -------------------image generation utils-----
@@ -24,7 +25,7 @@ class PngWriter:
             print(e)
         self.files_written.append([self.filepath,seed])
 
-    def unique_filename(self,seed,previouspath):
+    def unique_filename(self,seed,previouspath=None):
         revision = 1
 
         if previouspath is None:
@@ -61,6 +62,22 @@ class PngWriter:
         info = PngImagePlugin.PngInfo()
         info.add_text("Dream",prompt)
         image.save(path,"PNG",pnginfo=info)
+
+    def make_grid(self,image_list,rows=None,cols=None):
+        image_cnt = len(image_list)
+        if None in (rows,cols):
+            rows = floor(sqrt(image_cnt))  # try to make it square
+            cols = ceil(image_cnt/rows)
+        width  = image_list[0].width
+        height = image_list[0].height
+
+        grid_img = Image.new('RGB',(width*cols,height*rows))
+        for r in range(0,rows):
+            for c in range (0,cols):
+                i = r*rows + c
+                grid_img.paste(image_list[i],(c*width,r*height))
+
+        return grid_img
     
 class PromptFormatter():
     def __init__(self,t2i,opt):
@@ -80,8 +97,6 @@ class PromptFormatter():
         switches.append(f'-H{opt.height       or t2i.height}')
         switches.append(f'-C{opt.cfg_scale    or t2i.cfg_scale}')
         switches.append(f'-m{t2i.sampler_name}')
-        if opt.variants:
-            switches.append(f'-v{opt.variants}')
         if opt.init_img:
             switches.append(f'-I{opt.init_img}')
         if opt.strength and opt.init_img is not None:

ldm/simplet2i.py

@@ -20,7 +20,6 @@ from torch import autocast
 from contextlib import contextmanager, nullcontext
 import transformers
 import time
-import math
 import re
 import traceback
 

scripts/dream.py

@@ -153,54 +153,60 @@ def main_loop(t2i,outdir,parser,log,infile):
             continue
 
         normalized_prompt      = PromptFormatter(t2i,opt).normalize_prompt()
-        variants               = None
+        individual_images      = not opt.grid
 
         try:
             file_writer        = PngWriter(outdir,normalized_prompt,opt.batch_size)
-            callback           = file_writer.write_image
+            callback           = file_writer.write_image if individual_images else None
 
-            t2i.prompt2image(image_callback=callback,
+            image_list   = t2i.prompt2image(image_callback=callback,**vars(opt))
-                             **vars(opt))
+            results      = file_writer.files_written     if individual_images else image_list
-            results      = file_writer.files_written
 
-            if None not in (opt.variants,opt.init_img):
+            if opt.grid and len(results) > 0:
-                variants = generate_variants(t2i,outdir,opt,results)
+                grid_img = file_writer.make_grid([r[0] for r in results])
+                filename = file_writer.unique_filename(results[0][1])
+                seeds    = [a[1] for a in results]
+                results  = [[filename,seeds]]
+                metadata_prompt   = f'{normalized_prompt} -S{results[0][1]}'
+                file_writer.save_image_and_prompt_to_png(grid_img,metadata_prompt,filename)
 
         except AssertionError as e:
             print(e)
             continue
 
+        except OSError as e:
+            print(e)
+            continue
+
         print("Outputs:")
         write_log_message(t2i,normalized_prompt,results,log)
-        if variants is not None:
-            print('Variants:')
-            for vr in variants:
-                write_log_message(t2i,vr[0],vr[1],log)
 
     print("goodbye!")
 
-def generate_variants(t2i,outdir,opt,previous_gens):
+# variant generation is going to be superseded by a generalized
-    variants = []
+# "prompt-morph" functionality
-    print(f"Generating {opt.variants} variant(s)...")
+# def generate_variants(t2i,outdir,opt,previous_gens):
-    newopt = copy.deepcopy(opt)
+#     variants = []
-    newopt.iterations = 1
+#     print(f"Generating {opt.variants} variant(s)...")
-    newopt.variants   = None
+#     newopt = copy.deepcopy(opt)
-    for r in previous_gens:
+#     newopt.iterations = 1
-        newopt.init_img = r[0]
+#     newopt.variants   = None
-        prompt            = PromptFormatter(t2i,newopt).normalize_prompt()
+#     for r in previous_gens:
-        print(f"] generating variant for {newopt.init_img}")
+#         newopt.init_img = r[0]
-        for j in range(0,opt.variants):
+#         prompt            = PromptFormatter(t2i,newopt).normalize_prompt()
-            try:
+#         print(f"] generating variant for {newopt.init_img}")
-                file_writer        = PngWriter(outdir,prompt,newopt.batch_size)
+#         for j in range(0,opt.variants):
-                callback           = file_writer.write_image
+#             try:
-                t2i.prompt2image(image_callback=callback,**vars(newopt))
+#                 file_writer        = PngWriter(outdir,prompt,newopt.batch_size)
-                results           = file_writer.files_written
+#                 callback           = file_writer.write_image
-                variants.append([prompt,results])
+#                 t2i.prompt2image(image_callback=callback,**vars(newopt))
-            except AssertionError as e:
+#                 results           = file_writer.files_written
-                print(e)
+#                 variants.append([prompt,results])
-                continue
+#             except AssertionError as e:
-    print(f'{opt.variants} variants generated')
+#                 print(e)
-    return variants
+#                 continue
+#     print(f'{opt.variants} variants generated')
+#     return variants
                 
 
 def write_log_message(t2i,prompt,results,logfile):
@@ -209,9 +215,6 @@ def write_log_message(t2i,prompt,results,logfile):
     img_num    = 1
     seenit     = {}
 
-    seeds = [a[1] for a in results]
-    seeds = f"(seeds for individual images: {seeds})"
-
     for r in results:
         seed = r[1]
         log_message = (f'{r[0]}: {prompt} -S{seed}')
@@ -275,7 +278,8 @@ 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 for img2img mode (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="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
+# variants is going to be superseded by a generalized "prompt-morph" function
+#    parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
     parser.add_argument('-x','--skip_normalize',action='store_true',help="skip subprompt weight normalization")
     return parser
 

src/k-diffusion

@@ -1 +1 @@
-Subproject commit db5799068749bf3a6d5845120ed32df16b7d883b
+Subproject commit ef1bf07627c9a10ba9137e68a0206b844544a7d9