From 3d00251fc27fb9b5c67ec7bcdf2a6fb7ef497988 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Thu, 27 Jun 2024 19:37:00 -0500 Subject: [PATCH] modified prepare_mask_batch to work with sd3 --- adv_control/utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/adv_control/utils.py b/adv_control/utils.py index f16ac35..b8651f9 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -429,11 +429,15 @@ class manual_cast_clean_groupnorm(comfy.ops.manual_cast): # adapted from comfy/sample.py -def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim1=False): +def prepare_mask_batch(mask: Tensor, shape: Tensor, multiplier: int=1, match_dim1=False, match_shape=False): mask = mask.clone() - mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[2]*multiplier, shape[3]*multiplier), mode="bilinear") + mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*multiplier, shape[-1]*multiplier), mode="bilinear") if match_dim1: + if match_shape and len(shape) < 4: + raise Exception(f"match_dim1 cannot be True if shape is under 4 dims; was {len(shape)}.") mask = torch.cat([mask] * shape[1], dim=1) + if match_shape and len(shape) == 3: + mask = mask.squeeze(1) return mask @@ -823,10 +827,10 @@ class AdvancedControlBase: x[:] = x[:] * self.calc_latent_keyframe_mults(x=x, batched_number=batched_number) # apply masks, resizing mask to required dims if self.mask_cond_hint is not None: - masks = prepare_mask_batch(self.mask_cond_hint, x.shape) + masks = prepare_mask_batch(self.mask_cond_hint, x.shape, match_shape=True) x[:] = x[:] * masks if self.tk_mask_cond_hint is not None: - masks = prepare_mask_batch(self.tk_mask_cond_hint, x.shape) + masks = prepare_mask_batch(self.tk_mask_cond_hint, x.shape, match_shape=True) x[:] = x[:] * masks # apply timestep keyframe strengths if self._current_timestep_keyframe.strength != 1.0: