diff --git a/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py b/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py index da624b73..4a7390dd 100644 --- a/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py +++ b/extensions-builtin/forge_preprocessor_inpaint/scripts/preprocessor_inpaint.py @@ -10,6 +10,7 @@ import cv2 import einops import numpy as np import torch +import torch.nn.functional as F import yaml from omegaconf import OmegaConf @@ -55,7 +56,7 @@ class PreprocessorInpaintOnly(PreprocessorInpaint): latent_image = vae.encode(self.image.movedim(1, -1)) latent_image = process.sd_model.forge_objects.vae.first_stage_model.process_in(latent_image) - _, _, H, W = latent_image.shape + *_, H, W = latent_image.shape latent_mask = self.mask latent_mask = torch.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round() @@ -96,6 +97,10 @@ class PreprocessorInpaintOnly(PreprocessorInpaint): mask = torch.from_numpy(np.ascontiguousarray(mask).copy()).to(img).clip(0, 1) raw = self.image[0].to(img).clip(0, 1) img = img.clip(0, 1) + + if img.shape != mask.shape: # Interrupted + img = F.interpolate(img.unsqueeze(0), size=(mask.shape[1], mask.shape[2]), mode="nearest").squeeze(0) + new_results.append(raw * (1.0 - mask) + img * mask) a1111_batch_result.images = new_results diff --git a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py index f80082c1..e67852f3 100644 --- a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py +++ b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py @@ -20,13 +20,13 @@ class ControlLLLiteAnimaPatcher(ControlModelPatcher): return None _, metadata = load_torch_file(ckpt_path, return_metadata=True) inpaint_masked_input: bool = (metadata or {}).get("lllite.inpaint_masked_input", None) == "true" - return ControlLLLiteAnimaPatcher(state_dict, inpaint_masked_input=inpaint_masked_input) + return ControlLLLiteAnimaPatcher(state_dict, inpaint=inpaint_masked_input) - def __init__(self, state_dict, inpaint_masked_input: bool = False): + def __init__(self, state_dict: dict[str, torch.Tensor], inpaint: bool): super().__init__() self.state_dict = state_dict - self.inpaint_masked_input = inpaint_masked_input - self._lllite_net = None + self._is_inpaint = inpaint + self._lllite_net: ControlNetLLLiteDiT = None def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): unet = process.sd_model.forge_objects.unet @@ -38,21 +38,18 @@ class ControlLLLiteAnimaPatcher(ControlModelPatcher): self._lllite_net = ControlNetLLLiteDiT(dit, **cfg) load_lllite_weights_from_dict(self._lllite_net, self.state_dict) self._lllite_net = self._lllite_net.eval().to(device=device, dtype=dtype) + del self.state_dict cond_image = cond * 2.0 - 1.0 - if self._lllite_net.conditioning1.conv1.in_channels == 4: - b, c, h, w = cond_image.shape - if isinstance(mask, torch.Tensor): - inpaint_mask = mask.to(device=cond_image.device, dtype=cond_image.dtype) - if inpaint_mask.shape[-2:] != (h, w): - inpaint_mask = F.interpolate(inpaint_mask, size=(h, w), mode="nearest") - else: - inpaint_mask = torch.zeros(b, 1, h, w, device=cond_image.device, dtype=cond_image.dtype) - if self.inpaint_masked_input: - keep = (inpaint_mask < 0.5).to(cond_image.dtype) - cond_image = cond_image * keep + if self._is_inpaint: + assert isinstance(mask, torch.Tensor) + if mask.shape != cond_image.shape: + mask = F.interpolate(mask, size=(cond_image.shape[2], cond_image.shape[3]), mode="nearest") + inpaint_mask = mask.to(device=cond_image.device, dtype=cond_image.dtype) + cond_image = cond_image * (inpaint_mask < 0.5) inpaint_mask = inpaint_mask * 2.0 - 1.0 cond_image = torch.cat([cond_image, inpaint_mask], dim=1) + self._lllite_net.set_cond_image(cond_image.to(device=device, dtype=dtype)) self._lllite_net.set_multiplier(self.strength) self._lllite_net.set_step_range(num_steps=process.steps, start_percent=self.start_percent, end_percent=self.end_percent) diff --git a/modules_forge/shared.py b/modules_forge/shared.py index f02240cc..0bd5bda5 100644 --- a/modules_forge/shared.py +++ b/modules_forge/shared.py @@ -1,4 +1,9 @@ import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from modules_forge.supported_controlnet import ControlModelPatcher + from modules_forge.supported_preprocessor import Preprocessor from backend import utils from modules.paths_internal import models_path, normalized_filepath, parser @@ -34,19 +39,19 @@ os.makedirs(preprocessor_dir, exist_ok=True) diffusers_dir: str = os.path.join(models_path, "diffusers") os.makedirs(diffusers_dir, exist_ok=True) -supported_preprocessors = {} -supported_control_models = [] +supported_preprocessors: dict[str, "Preprocessor"] = {} +supported_control_models: list["ControlModelPatcher"] = [] -def add_supported_preprocessor(preprocessor): +def add_supported_preprocessor(preprocessor: "Preprocessor"): supported_preprocessors[preprocessor.name] = preprocessor -def add_supported_control_model(control_model): +def add_supported_control_model(control_model: "ControlModelPatcher"): supported_control_models.append(control_model) -def try_load_supported_control_model(ckpt_path): +def try_load_supported_control_model(ckpt_path: os.PathLike): state_dict = utils.load_torch_file(ckpt_path, safe_load=True) for supported_type in supported_control_models: state_dict_copy = {k: v for k, v in state_dict.items()}