From 58d06a9b736387d0deccb940e449275ac45be999 Mon Sep 17 00:00:00 2001
From: EllangoK <karun.ellango7@gmail.com>
Date: Tue, 31 Jan 2023 02:44:07 -0500
Subject: [PATCH] draw_inner_margin now works with legends padding

---
 modules/images.py   | 33 ++++++++++++++-------------------
 scripts/xyz_grid.py |  7 ++++---
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/modules/images.py b/modules/images.py
index c5fba30e..dbf2cb19 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -216,28 +216,23 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
 
         draw_texts(d, x, y, ver_texts[row], fnt, fontsize)
 
-    return result
-
-def draw_inner_margins(im, margin, rows, cols, margin_color=(255, 255, 255)):
-    cell_width = im.width // cols
-    cell_height = im.height // rows
-
-    padded_width = im.width + (cols - 1) * margin
-    padded_height = im.height + (rows - 1) * margin
-
-    padded_im = Image.new("RGB", (padded_width, padded_height), margin_color)
+    return result, pad_left, pad_top
 
+def draw_inner_margins(im, margin, rows, cols, pad_left, pad_top, margin_color=(255, 255, 255)):
+    if margin == 0: # Early exit, not strictly necessary
+        return im
+    cell_width = (im.width - pad_left) // cols
+    cell_height = (im.height - pad_top) // rows
+    padded_im = Image.new("RGB", (im.width + (cols - 1) * margin, im.height + (rows - 1) * margin), margin_color)
     for r in range(rows):
         for c in range(cols):
-            cell_location = (
-                c * cell_width,
-                r * cell_height,
-                (c + 1) * cell_width,
-                (r + 1) * cell_height
-            )
-            paste_location = (c * (cell_width + margin), r * (cell_height + margin))
-            padded_im.paste(im.crop(cell_location), paste_location)
-
+            cell_location = (c * cell_width + pad_left * (c != 0), 
+                             r * cell_height + pad_top * (r != 0),
+                             (c + 1) * cell_width + pad_left, 
+                             (r + 1) * cell_height + pad_top)
+            paste_location = (c * (cell_width + margin) + pad_left * (c != 0), 
+                              r * (cell_height + margin) + pad_top * (r != 0))
+            padded_im.paste(im.crop(cell_location), paste_location) 
     return padded_im
 
 def draw_prompt_matrix(im, width, height, all_prompts):
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index ff9fc356..f2c36b5f 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -291,9 +291,10 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
         start_index = i * len(xs) * len(ys)
         end_index = start_index + len(xs) * len(ys)
         grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys))
-        grid = images.draw_inner_margins(grid, margin_size, len(ys), len(xs))
+        pad_left, pad_top = 0, 0
         if draw_legend:
-            grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+            grid, pad_left, pad_top = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+        grid = images.draw_inner_margins(grid, margin_size, len(ys), len(xs), pad_left, pad_top)
         sub_grids[i] = grid
         if include_sub_grids and len(zs) > 1:
             processed_result.images.insert(i+1, grid)
@@ -301,7 +302,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
     sub_grid_size = sub_grids[0].size
     z_grid = images.image_grid(sub_grids, rows=1)
     if draw_legend:
-        z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])
+        z_grid, _, _ = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])
     processed_result.images[0] = z_grid
 
     return processed_result, sub_grids