add batch

2025-04-02 20:32:49 +08:00
parent b0406e4adc
commit 0915d2cd2e
1 changed files with 4 additions and 0 deletions
@@ -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()