touch up

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

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)

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()}