move make_grid into image_utils

ldm/dream/image_util.py

@@ -1,3 +1,4 @@
+from math import sqrt, floor, ceil
 from PIL import Image
 
 class InitImageResizer():
@@ -51,4 +52,22 @@ class InitImageResizer():
 
         return new_image
 
+def make_grid(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))
+    i = 0
+    for r in range(0, rows):
+        for c in range(0, cols):
+            if i >= len(image_list):
+                break
+            grid_img.paste(image_list[i], (c * width, r * height))
+            i = i + 1
+
+    return grid_img
 

ldm/dream/pngwriter.py

@@ -2,16 +2,13 @@
 Two helper classes for dealing with PNG images and their path names.
 PngWriter -- Converts Images generated by T2I into PNGs, finds
              appropriate names for them, and writes prompt metadata
-             into the PNG. Intended to be subclassable in order to
+             into the PNG.
-             create more complex naming schemes, including using the
-             prompt for file/directory names.
 PromptFormatter -- Utility for converting a Namespace of prompt parameters
              back into a formatted prompt string with command-line switches.
 """
 import os
 import re
-from math import sqrt, floor, ceil
+from PIL import PngImagePlugin
-from PIL import Image, PngImagePlugin
 
 # -------------------image generation utils-----
 
@@ -42,26 +39,6 @@ class PngWriter:
         image.save(path, 'PNG', pnginfo=info)
         return path
 
-    # TODO move this to its own helper function; it's not really a method of pngwriter
-    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))
-        i = 0
-        for r in range(0, rows):
-            for c in range(0, cols):
-                if i>=len(image_list):
-                    break
-                grid_img.paste(image_list[i], (c * width, r * height))
-                i = i + 1
-
-        return grid_img
-
 
 class PromptFormatter:
     def __init__(self, t2i, opt):

scripts/dream.py

@@ -12,6 +12,7 @@ import time
 import ldm.dream.readline
 from ldm.dream.pngwriter import PngWriter, PromptFormatter
 from ldm.dream.server import DreamServer, ThreadingDreamServer
+from ldm.dream.image_util import make_grid
 
 def main():
     """Initialize command-line parsers and the diffusion model"""
@@ -226,7 +227,7 @@ def main_loop(t2i, outdir, prompt_as_dir, parser, infile):
             t2i.prompt2image(image_callback=image_writer, **vars(opt))
 
             if do_grid and len(grid_images) > 0:
-                grid_img = file_writer.make_grid(list(grid_images.values()))
+                grid_img = make_grid(list(grid_images.values()))
                 first_seed = next(iter(seeds))
                 filename = f'{prefix}.{first_seed}.png'
                 # TODO better metadata for grid images