From e672cfb07418a1a3130d3bf21c14a0d3819f81fb Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 1 Jan 2023 18:37:37 +0300
Subject: rework of callback for #6094

---
 modules/images.py           | 10 ++++++----
 modules/script_callbacks.py | 26 +++++++++++++++-----------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/modules/images.py b/modules/images.py
index 719aaf3b..f84fd485 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -39,12 +39,14 @@ def image_grid(imgs, batch_size=1, rows=None):
 
     cols = math.ceil(len(imgs) / rows)
 
+    params = script_callbacks.ImageGridLoopParams(imgs, cols, rows)
+    script_callbacks.image_grid_callback(params)
+
     w, h = imgs[0].size
-    grid = Image.new('RGB', size=(cols * w, rows * h), color='black')
+    grid = Image.new('RGB', size=(params.cols * w, params.rows * h), color='black')
 
-    for i, img in enumerate(imgs):
-        script_callbacks.image_grid_loop_callback(img)
-        grid.paste(img, box=(i % cols * w, i // cols * h))
+    for i, img in enumerate(params.imgs):
+        grid.paste(img, box=(i % params.cols * w, i // params.cols * h))
 
     return grid
 
diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 0c854407..de69fd9f 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -52,8 +52,10 @@ class UiTrainTabParams:
 
 
 class ImageGridLoopParams:
-    def __init__(self, img):
-        self.img = img
+    def __init__(self, imgs, cols, rows):
+        self.imgs = imgs
+        self.cols = cols
+        self.rows = rows
 
 
 ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
@@ -68,7 +70,7 @@ callback_map = dict(
     callbacks_cfg_denoiser=[],
     callbacks_before_component=[],
     callbacks_after_component=[],
-    callbacks_image_grid_loop=[],
+    callbacks_image_grid=[],
 )
 
 
@@ -160,12 +162,14 @@ def after_component_callback(component, **kwargs):
         except Exception:
             report_exception(c, 'after_component_callback')
 
-def image_grid_loop_callback(component, **kwargs):
-    for c in callback_map['callbacks_image_grid_loop']:
+
+def image_grid_callback(params: ImageGridLoopParams):
+    for c in callback_map['callbacks_image_grid']:
         try:
-            c.callback(component, **kwargs)
+            c.callback(params)
         except Exception:
-            report_exception(c, 'image_grid_loop')
+            report_exception(c, 'image_grid')
+
 
 def add_callback(callbacks, fun):
     stack = [x for x in inspect.stack() if x.filename != __file__]
@@ -269,9 +273,9 @@ def on_after_component(callback):
     add_callback(callback_map['callbacks_after_component'], callback)
 
 
-def on_image_grid_loop(callback):
-    """register a function to be called inside the image grid loop.
+def on_image_grid(callback):
+    """register a function to be called before making an image grid.
     The callback is called with one argument:
-       - params: ImageGridLoopParams - parameters to be used inside the image grid loop.
+       - params: ImageGridLoopParams - parameters to be used for grid creation. Can be modified.
     """
-    add_callback(callback_map['callbacks_image_grid_loop'], callback)
+    add_callback(callback_map['callbacks_image_grid'], callback)
-- 
cgit v1.2.1