From 4281f255d5e7c67515d619f53654be59a6fc1e13 Mon Sep 17 00:00:00 2001
From: wywywywy <wywywywy@gmail.com>
Date: Thu, 20 Oct 2022 15:31:09 +0100
Subject: [PATCH] Implemented batch count logic to Outpainting mk2

---
 scripts/outpainting_mk_2.py | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/scripts/outpainting_mk_2.py b/scripts/outpainting_mk_2.py
index a6468e09..02e655e9 100644
--- a/scripts/outpainting_mk_2.py
+++ b/scripts/outpainting_mk_2.py
@@ -242,21 +242,37 @@ class Script(scripts.Script):
             out = out.crop((0, 0, res_w, res_h))
             return out
 
-        img = init_image
+        batch_count = p.n_iter
+        p.n_iter = 1
+        state.job_count = batch_count
+        all_images = []
 
-        if left > 0:
-            img = expand(img, left, is_left=True)
-        if right > 0:
-            img = expand(img, right, is_right=True)
-        if up > 0:
-            img = expand(img, up, is_top=True)
-        if down > 0:
-            img = expand(img, down, is_bottom=True)
+        for i in range(batch_count):
+            img = init_image
+            state.job = f"Batch {i + 1} out of {state.job_count}"
 
-        res = Processed(p, [img], initial_seed_and_info[0], initial_seed_and_info[1])
+            if left > 0:
+                img = expand(img, left, is_left=True)
+            if right > 0:
+                img = expand(img, right, is_right=True)
+            if up > 0:
+                img = expand(img, up, is_top=True)
+            if down > 0:
+                img = expand(img, down, is_bottom=True)
+
+            all_images.append(img)
+
+        combined_grid_image = images.image_grid(all_images)
+        if opts.return_grid:
+            all_images = [combined_grid_image] + all_images
+        
+        res = Processed(p, all_images, initial_seed_and_info[0], initial_seed_and_info[1])
 
         if opts.samples_save:
             images.save_image(img, p.outpath_samples, "", res.seed, p.prompt, opts.grid_format, info=res.info, p=p)
 
+        if opts.grid_save:
+            images.save_image(combined_grid_image, p.outpath_grids, "grid", res.seed, p.prompt, opts.grid_format, info=res.info, short_filename=not opts.grid_extended_filename, grid=True, p=p)
+
         return res