From b35357b69abaedbddf5f03dfa28d03b504b654f9 Mon Sep 17 00:00:00 2001
From: nullquant <merlinkun@yandex.ru>
Date: Sat, 4 May 2024 17:55:29 +0300
Subject: [PATCH] latent batch and size complete rewrite

---
 brushnet/brushnet.py | 12 --------
 brushnet_nodes.py    | 68 ++++++++++++++++++++++++++++++--------------
 2 files changed, 47 insertions(+), 33 deletions(-)

diff --git a/brushnet/brushnet.py b/brushnet/brushnet.py
index d93c49a..39003c0 100644
--- a/brushnet/brushnet.py
+++ b/brushnet/brushnet.py
@@ -822,18 +822,6 @@ class BrushNetModel(ModelMixin, ConfigMixin):
         
         # 2. pre-process
 
-        if sample.shape[2] != brushnet_cond.shape[2] or sample.shape[3] != brushnet_cond.shape[3]:
-            print("Image "+ str(brushnet_cond.shape) +" and latent " + str(sample.shape) + " have different size, resizing image")
-            brushnet_cond = torch.nn.functional.interpolate(
-                brushnet_cond, size=(
-                    sample.shape[2], 
-                    sample.shape[3],
-                ), mode='bicubic',
-            ).to(sample.device).to(sample.dtype)
-
-        if sample.shape[0] > 1:
-            brushnet_cond = torch.cat([brushnet_cond] * sample.shape[0], dim=0).to(sample.device)
-
         brushnet_cond=torch.concat([sample,brushnet_cond],1)
         sample = self.conv_in_condition(brushnet_cond)
 
diff --git a/brushnet_nodes.py b/brushnet_nodes.py
index d111b3b..01971c2 100644
--- a/brushnet_nodes.py
+++ b/brushnet_nodes.py
@@ -408,7 +408,44 @@ def brushnet_inference(x, timesteps, transformer_options):
         added_cond_kwargs2 = mp['brushnet_add_embeds2']
         step = mp['step']
         total_steps = mp['total_steps']
-        
+
+        check = False
+        if (x.shape[0], x.shape[2], x.shape[3]) != (conditioning_latents.shape[0], conditioning_latents.shape[2], conditioning_latents.shape[3]):
+            if step == 0:
+                print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape)
+            check = True
+
+        if x.shape[0] == 2:
+            if step == 0:
+                print('BrushNet inference: do_classifier_free_guidance')
+            conditioning_latents = conditioning_latents2
+            prompt_embeds = prompt_embeds2
+            added_cond_kwargs = added_cond_kwargs2
+        elif x.shape[0] > 2:
+            if step == 0:
+                print('BrushNet inference: latent batch ' + str(x.shape[0]))
+            if x.shape[0] % 2 == 0:
+                if step == 0:
+                    print('BrushNet inference: do_classifier_free_guidance')
+                conditioning_latents = torch.cat([conditioning_latents2] * (x.shape[0] // 2), dim=0).to(x.device)
+                prompt_embeds = prompt_embeds2
+                added_cond_kwargs = added_cond_kwargs2
+            else:
+                conditioning_latents = torch.cat([conditioning_latents] * x.shape[0], dim=0).to(x.device)
+
+        if x.shape[2] != conditioning_latents.shape[2] or x.shape[3] != conditioning_latents.shape[3]:
+            if step == 0:
+                print('BrushNet inference: image '+ str(conditioning_latents.shape) + ' and latent ' + str(x.shape) + ' have different size, resizing image')
+            conditioning_latents = torch.nn.functional.interpolate(
+                conditioning_latents, size=(
+                    x.shape[2], 
+                    x.shape[3],
+                ), mode='bicubic',
+            ).to(x.device).to(x.dtype)
+
+        if check and step == 0:
+            print('BrushNet inference: sample', x.shape, ', CL', conditioning_latents.shape)
+
         brushnet_keep = []
         for i in range(total_steps):
             keeps = [
@@ -419,26 +456,15 @@ def brushnet_inference(x, timesteps, transformer_options):
 
         cond_scale = brushnet_conditioning_scale * brushnet_keep[step]
 
-        if x.shape[0] == 2:
-            return brushnet(x.to(brushnet.device),
-                            encoder_hidden_states=prompt_embeds2,
-                            brushnet_cond=conditioning_latents2,
-                            timestep = timesteps.to(brushnet.device),
-                            conditioning_scale=cond_scale,
-                            guess_mode=False,
-                            added_cond_kwargs=added_cond_kwargs2,
-                            return_dict=False,
-                        )
-        else:
-            return brushnet(x.to(brushnet.device),
-                            encoder_hidden_states=prompt_embeds,
-                            brushnet_cond=conditioning_latents,
-                            timestep = timesteps.to(brushnet.device),
-                            conditioning_scale=cond_scale,
-                            guess_mode=False,
-                            added_cond_kwargs=added_cond_kwargs,
-                            return_dict=False,
-                        )
+        return brushnet(x.to(brushnet.device),
+                        encoder_hidden_states=prompt_embeds,
+                        brushnet_cond=conditioning_latents,
+                        timestep = timesteps.to(brushnet.device),
+                        conditioning_scale=cond_scale,
+                        guess_mode=False,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )
     else:
         print('BrushNet model is not a BrushNetModel class')
         return ([], 0, [])