Fix for batch size > 1 with CNet

2024-05-29 23:25:53 +03:30 · 2024-05-29 23:25:53 +03:30 · 5a9ec7fbc7
parent cdacb3d8f1
commit 5a9ec7fbc7
1 changed files with 2 additions and 0 deletions
--- a/scripts/npw.py
+++ b/scripts/npw.py
@ -85,6 +85,8 @@ class Script(scripts.Script):

    def denoiser_callback(self, params):
        def concat_and_lerp(empty, tensor, weight):
+            if empty.shape[0] != tensor.shape[0]:
+                empty = empty.expand(tensor.shape[0], *empty.shape[1:])
            if tensor.shape[1] > empty.shape[1]:
                num_concatenations = tensor.shape[1] // empty.shape[1]
                empty_concat = torch.cat([empty] * num_concatenations, dim=1)