From 8215a3c6df759d5b47d62ac3c24e9de86333603e Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sun, 30 Jul 2023 00:14:36 +0200
Subject: [PATCH] Add batch override

Also fixes image combine, might need to re-add the nodes
---
 nodes/misc.py           | 2 +-
 nodes/remote_control.py | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/nodes/misc.py b/nodes/misc.py
index 6cd94f3..6e74af7 100644
--- a/nodes/misc.py
+++ b/nodes/misc.py
@@ -21,7 +21,7 @@ class CombineImageBatch:
 
 	def combine_images(self,images_a,images_b):
 		try:
-			out = torch.cat((images_a,images_b))
+			out = torch.cat((images_a,images_b), 0)
 		except RuntimeError:
 			print(f"Imagine size mismatch! {images_a.size()}, {images_b.size()}")
 			out = images_a
diff --git a/nodes/remote_control.py b/nodes/remote_control.py
index 5dcb9d9..bfd15de 100644
--- a/nodes/remote_control.py
+++ b/nodes/remote_control.py
@@ -80,6 +80,7 @@ class QueueRemote:
 					"default": "http://127.0.0.1:8188/",
 				}),
 				"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+				"batch_override": ("INT", {"default": 0, "min": 0, "max": 8}),
 				"system": (["windows", "posix"],),
 				"enabled": (["true", "false", "remote"],{"default": "true"}),
 				"node_id": ("INT", {"default": 1, "min": 1, "max": 64}),
@@ -94,14 +95,18 @@ class QueueRemote:
 	FUNCTION = "queue_on_remote"
 	CATEGORY = "remote"
 	
-	def queue_on_remote(self, remote_url, seed, system, enabled, node_id, prompt):
+	def queue_on_remote(self, remote_url, seed, batch_override, system, enabled, node_id, prompt):
 		def get_max_batch_size(prompt):
 			bs = 1
 			for node in prompt.keys():
 				if prompt[node]["class_type"] == "EmptyLatentImage":
 					for k,v in prompt[node]["inputs"].items():
 						if k == "batch_size":
-							bs = max(bs,int(v))
+							if batch_override > 0:
+								bs = batch_override
+								prompt[node]["inputs"]["batch_size"] = batch_override
+							else:
+								bs = max(bs,int(v))
 			return bs
 
 		if enabled == "false":