From 2aee7ea2049218c01de007d5a3faec34a5cec908 Mon Sep 17 00:00:00 2001 From: nullquant Date: Sat, 4 May 2024 13:10:04 +0300 Subject: [PATCH] latent batch --- brushnet/brushnet.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/brushnet/brushnet.py b/brushnet/brushnet.py index 93a32bd..d93c49a 100644 --- a/brushnet/brushnet.py +++ b/brushnet/brushnet.py @@ -831,6 +831,9 @@ class BrushNetModel(ModelMixin, ConfigMixin): ), 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)