From 0915d2cd2eb7ca22c64a53fd9e9a74f0d79d6323 Mon Sep 17 00:00:00 2001 From: jax Date: Wed, 2 Apr 2025 20:32:49 +0800 Subject: [PATCH] add batch --- nodes/comfy_nodes.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nodes/comfy_nodes.py b/nodes/comfy_nodes.py index e96d855..08ffc5c 100644 --- a/nodes/comfy_nodes.py +++ b/nodes/comfy_nodes.py @@ -207,6 +207,10 @@ class EasyControlGenerate: # Convert numpy array to torch.Tensor image = torch.from_numpy(image).float() + # Add batch dimension to make it [batch, height, width, channels] + if image.dim() == 3: # [height, width, channels] + image = image.unsqueeze(0) # Add batch dimension to make it [1, height, width, channels] + # Clear cache after generation for name, attn_processor in transformer.attn_processors.items(): attn_processor.bank_kv.clear()