CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

mask

Browse Source
This commit is contained in:
Haoming 2026-05-25 15:37:49 +08:00
parent 11db56306c
commit 34fd1bb415
1 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py
View File

@ -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