Prompt matrix now draws text like in demo.

webui.py

@@ -1,11 +1,10 @@
-import PIL
 import argparse, os, sys, glob
 import torch
 import torch.nn as nn
 import numpy as np
 import gradio as gr
 from omegaconf import OmegaConf
-from PIL import Image
+from PIL import Image, ImageFont, ImageDraw
 from itertools import islice
 from einops import rearrange, repeat
 from torch import autocast
@@ -76,23 +75,6 @@ def load_model_from_config(config, ckpt, verbose=False):
     return model
 
 
-def load_img_pil(img_pil):
-    image = img_pil.convert("RGB")
-    w, h = image.size
-    print(f"loaded input image of size ({w}, {h})")
-    w, h = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 64
-    image = image.resize((w, h), resample=PIL.Image.LANCZOS)
-    print(f"cropped image to size ({w}, {h})")
-    image = np.array(image).astype(np.float32) / 255.0
-    image = image[None].transpose(0, 3, 1, 2)
-    image = torch.from_numpy(image)
-    return 2. * image - 1.
-
-
-def load_img(path):
-    return load_img_pil(Image.open(path))
-
-
 class CFGDenoiser(nn.Module):
     def __init__(self, model):
         super().__init__()
@@ -179,6 +161,71 @@ def image_grid(imgs, batch_size, round_down=False):
     return grid
 
 
+def draw_prompt_matrix(im, width, height, all_prompts):
+    def wrap(text, d, font, line_length):
+        lines = ['']
+        for word in text.split():
+            line = f'{lines[-1]} {word}'.strip()
+            if d.textlength(line, font=font) <= line_length:
+                lines[-1] = line
+            else:
+                lines.append(word)
+        return '\n'.join(lines)
+
+    def draw_texts(pos, x, y, texts, sizes):
+        for i, (text, size) in enumerate(zip(texts, sizes)):
+            active = pos & (1 << i) != 0
+
+            if not active:
+                text = '\u0336'.join(text) + '\u0336'
+
+            d.multiline_text((x, y + size[1] / 2), text, font=fnt, fill=color_active if active else color_inactive, anchor="mm", align="center")
+
+            y += size[1] + line_spacing
+
+    fontsize = (width + height) // 25
+    line_spacing = fontsize // 2
+    fnt = ImageFont.truetype("arial.ttf", fontsize)
+    color_active = (0, 0, 0)
+    color_inactive = (153, 153, 153)
+
+    pad_top = height // 4
+    pad_left = width * 3 // 4
+
+    cols = im.width // width
+    rows = im.height // height
+
+    prompts = all_prompts[1:]
+
+    result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white")
+    result.paste(im, (pad_left, pad_top))
+
+    d = ImageDraw.Draw(result)
+
+    boundary = math.ceil(len(prompts) / 2)
+    prompts_horiz = [wrap(x, d, fnt, width) for x in prompts[:boundary]]
+    prompts_vert = [wrap(x, d, fnt, pad_left) for x in prompts[boundary:]]
+
+    sizes_hor = [(x[2] - x[0], x[3] - x[1]) for x in [d.multiline_textbbox((0, 0), x, font=fnt) for x in prompts_horiz]]
+    sizes_ver = [(x[2] - x[0], x[3] - x[1]) for x in [d.multiline_textbbox((0, 0), x, font=fnt) for x in prompts_vert]]
+    hor_text_height = sum([x[1] + line_spacing for x in sizes_hor]) - line_spacing
+    ver_text_height = sum([x[1] + line_spacing for x in sizes_ver]) - line_spacing
+
+    for col in range(cols):
+        x = pad_left + width * col + width / 2
+        y = pad_top / 2 - hor_text_height / 2
+
+        draw_texts(col, x, y, prompts_horiz, sizes_hor)
+
+    for row in range(rows):
+        x = pad_left / 2
+        y = pad_top + height * row + height / 2 - ver_text_height / 2
+
+        draw_texts(row, x, y, prompts_vert, sizes_ver)
+
+    return result
+
+
 def dream(prompt: str, ddim_steps: int, sampler_name: str, use_GFPGAN: bool, prompt_matrix: bool, ddim_eta: float, n_iter: int, n_samples: int, cfg_scale: float, seed: int, height: int, width: int):
     torch.cuda.empty_cache()
 
@@ -212,30 +259,23 @@ def dream(prompt: str, ddim_steps: int, sampler_name: str, use_GFPGAN: bool, pro
     grid_count = len(os.listdir(outpath)) - 1
 
     prompt_matrix_prompts = []
-    comment = ""
+    prompt_matrix_parts = []
     if prompt_matrix:
         keep_same_seed = True
-        comment = "Image prompts:\n\n"
 
-        items = prompt.split("|")
-        combination_count = 2 ** (len(items)-1)
+        prompt_matrix_parts = prompt.split("|")
+        combination_count = 2 ** (len(prompt_matrix_parts)-1)
         for combination_num in range(combination_count):
-            current = items[0]
-            label = 'A'
+            current = prompt_matrix_parts[0]
 
-            for n, text in enumerate(items[1:]):
+            for n, text in enumerate(prompt_matrix_parts[1:]):
                 if combination_num & (2**n) > 0:
                     current += ("" if text.strip().startswith(",") else ", ") + text
-                    label += chr(ord('B') + n)
-
-            comment += " - " + label + "\n"
 
             prompt_matrix_prompts.append(current)
         n_iter = math.ceil(len(prompt_matrix_prompts) / batch_size)
 
-        comment += "\nwhere:\n"
-        for n, text in enumerate(items):
-            comment += "  " + chr(ord('A') + n) + " = " + items[n] + "\n"
+        print(f"Prompt matrix will create {len(prompt_matrix_prompts)} images using a total of {n_iter} batches.")
 
     precision_scope = autocast if opt.precision == "autocast" else nullcontext
     output_images = []
@@ -262,7 +302,7 @@ def dream(prompt: str, ddim_steps: int, sampler_name: str, use_GFPGAN: bool, pro
             x_samples_ddim = model.decode_first_stage(samples_ddim)
             x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
 
-            if not opt.skip_save or not opt.skip_grid:
+            if prompt_matrix or not opt.skip_save or not opt.skip_grid:
                 for i, x_sample in enumerate(x_samples_ddim):
                     x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                     x_sample = x_sample.astype(np.uint8)
@@ -279,24 +319,23 @@ def dream(prompt: str, ddim_steps: int, sampler_name: str, use_GFPGAN: bool, pro
                     output_images.append(image)
                     base_count += 1
 
-        if not opt.skip_grid:
-            # additionally, save as grid
+        if prompt_matrix or not opt.skip_grid:
             grid = image_grid(output_images, batch_size, round_down=prompt_matrix)
+
+            if prompt_matrix:
+                grid = draw_prompt_matrix(grid, width, height, prompt_matrix_parts)
+                output_images.insert(0, grid)
+
             grid.save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
             grid_count += 1
 
-
-    if sampler is not None:
-        del sampler
+    del sampler
 
     info = f"""
 {prompt}
 Steps: {ddim_steps}, Sampler: {sampler_name}, CFG scale: {cfg_scale}, Seed: {seed}{', GFPGAN' if use_GFPGAN and GFPGAN is not None else ''}
     """.strip()
 
-    if len(comment) > 0:
-        info += "\n\n" + comment
-
     return output_images, seed, info
 
 class Flagging(gr.FlaggingCallback):
@@ -350,7 +389,7 @@ dream_interface = gr.Interface(
         gr.Checkbox(label='Create prompt matrix (separate multiple prompts using |, and get all combinations of them)', value=False),
         gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="DDIM ETA", value=0.0, visible=False),
         gr.Slider(minimum=1, maximum=16, step=1, label='Batch count (how many batches of images to generate)', value=1),
-        gr.Slider(minimum=1, maximum=4, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
+        gr.Slider(minimum=1, maximum=8, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
         gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', value=7.0),
         gr.Number(label='Seed', value=-1),
         gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512),
@@ -389,7 +428,7 @@ def translation(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, ddim_e
     grid_count = len(os.listdir(outpath)) - 1
 
     image = init_img.convert("RGB")
-    image = image.resize((width, height), resample=PIL.Image.Resampling.LANCZOS)
+    image = image.resize((width, height), resample=Image.Resampling.LANCZOS)
     image = np.array(image).astype(np.float32) / 255.0
     image = image[None].transpose(0, 3, 1, 2)
     image = torch.from_numpy(image)
@@ -466,7 +505,7 @@ img2img_interface = gr.Interface(
         gr.Checkbox(label='Fix faces using GFPGAN', value=False, visible=GFPGAN is not None),
         gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="DDIM ETA", value=0.0, visible=False),
         gr.Slider(minimum=1, maximum=16, step=1, label='Batch count (how many batches of images to generate)', value=1),
-        gr.Slider(minimum=1, maximum=4, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
+        gr.Slider(minimum=1, maximum=8, step=1, label='Batch size (how many images are in a batch; memory-hungry)', value=1),
         gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='Classifier Free Guidance Scale (how strongly the image should follow the prompt)', value=7.0),
         gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising Strength', value=0.75),
         gr.Number(label='Seed', value=-1),
@@ -494,7 +533,7 @@ def run_GFPGAN(image, strength):
     res = Image.fromarray(restored_img)
 
     if strength < 1.0:
-        res = PIL.Image.blend(image, res, strength)
+        res = Image.blend(image, res, strength)
 
     return res