From 34fd1bb415ecb03a8840ce366fa151da11a2c08e Mon Sep 17 00:00:00 2001 From: Haoming Date: Mon, 25 May 2026 15:37:49 +0800 Subject: [PATCH] mask --- .../scripts/forge_controllllite.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py index e67852f3..6f8a7dc3 100644 --- a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py +++ b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py @@ -28,6 +28,12 @@ class ControlLLLiteAnimaPatcher(ControlModelPatcher): self._is_inpaint = inpaint self._lllite_net: ControlNetLLLiteDiT = None + @staticmethod + def _is_black_on_white(image: torch.Tensor) -> bool: + grey = image.mean(dim=1).squeeze(0) + white_ratio = (grey > 0.9).float().mean().item() + return white_ratio > 0.85 + def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): unet = process.sd_model.forge_objects.unet device, dtype = unet.load_device, unet.model.computation_dtype @@ -40,6 +46,13 @@ class ControlLLLiteAnimaPatcher(ControlModelPatcher): self._lllite_net = self._lllite_net.eval().to(device=device, dtype=dtype) del self.state_dict + if (not self._is_inpaint) and (mask is not None): + if inv := self._is_black_on_white(cond): + cond = 1.0 - cond + cond *= mask + if inv: + cond = 1.0 - cond + cond_image = cond * 2.0 - 1.0 if self._is_inpaint: assert isinstance(mask, torch.Tensor) @@ -74,6 +87,9 @@ class ControlLLLitePatcher(ControlModelPatcher): def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): unet = process.sd_model.forge_objects.unet + if mask is not None: + cond *= mask + unet = LLLiteLoader.load_lllite(model=unet, state_dict=self.state_dict, cond_image=cond.movedim(1, -1), strength=self.strength, steps=process.steps, start_percent=self.start_percent, end_percent=self.end_percent) process.sd_model.forge_objects.unet = unet