diff --git a/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite.py b/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite.py index 3bf7760d..9416ce8b 100644 --- a/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite.py +++ b/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite.py @@ -1,17 +1,21 @@ +# https://github.com/kohya-ss/ControlNet-LLLite-ComfyUI/blob/main/node_control_net_lllite.py + +import logging import math +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from backend.patcher.unet import UnetPatcher + import torch +import torch.nn as nn + +from backend.state_dict import load_state_dict + +logger = logging.getLogger("ControlNet") -def extra_options_to_module_prefix(extra_options): - # extra_options = {'transformer_index': 2, 'block_index': 8, 'original_shape': [2, 4, 128, 128], 'block': ('input', 7), 'n_heads': 20, 'dim_head': 64} - - # block is: [('input', 4), ('input', 5), ('input', 7), ('input', 8), ('middle', 0), - # ('output', 0), ('output', 1), ('output', 2), ('output', 3), ('output', 4), ('output', 5)] - # transformer_index is: [0, 1, 2, 3, 4, 5, 6, 7, 8], for each block - # block_index is: 0-1 or 0-9, depends on the block - # input 7 and 8, middle has 10 blocks - - # make module name from extra_options +def extra_options_to_module_prefix(extra_options: dict) -> str: block = extra_options["block"] block_index = extra_options["block_index"] if block[0] == "input": @@ -21,16 +25,14 @@ def extra_options_to_module_prefix(extra_options): elif block[0] == "output": module_pfx = f"lllite_unet_output_blocks_{block[1]}_1_transformer_blocks_{block_index}" else: - raise Exception("invalid block name") + raise ValueError("invalid block name") return module_pfx -def load_control_net_lllite_patch(ctrl_sd, cond_image, multiplier, num_steps, start_percent, end_percent, model_dtype): - # calculate start and end step +def load_control_net_lllite_patch(ctrl_sd: dict, cond_image: torch.Tensor, multiplier: float, num_steps: int, start_percent: float, end_percent: float, *, model_dtype: torch.dtype): start_step = math.floor(num_steps * start_percent) if start_percent > 0 else 0 end_step = math.floor(num_steps * end_percent) if end_percent > 0 else num_steps - # split each weights for each module module_weights = {} for key, value in ctrl_sd.items(): fragments = key.split(".") @@ -41,10 +43,8 @@ def load_control_net_lllite_patch(ctrl_sd, cond_image, multiplier, num_steps, st module_weights[module_name] = {} module_weights[module_name][weight_name] = value - # load each module modules = {} for module_name, weights in module_weights.items(): - # ここの自動判定を何とかしたい if "conditioning1.4.weight" in weights: depth = 3 elif weights["conditioning1.2.weight"].shape[-1] == 4: @@ -63,30 +63,28 @@ def load_control_net_lllite_patch(ctrl_sd, cond_image, multiplier, num_steps, st num_steps=num_steps, start_step=start_step, end_step=end_step, - dtype=model_dtype ) - info = module.load_state_dict(weights) - modules[module_name] = module + load_state_dict(module, weights) + modules[module_name] = module.eval().to(dtype=model_dtype) if len(modules) == 1: module.is_first = True - print(f"{len(modules)} modules") + logger.info(f"Loaded Control-LLLite ({len(modules)} modules)") - # cond imageをセットする - cond_image = cond_image.permute(0, 3, 1, 2) # b,h,w,3 -> b,3,h,w - cond_image = cond_image * 2.0 - 1.0 # 0-1 -> -1-+1 + cond_image = cond_image.permute(0, 3, 1, 2) + cond_image = cond_image * 2.0 - 1.0 for module in modules.values(): module.set_cond_image(cond_image) class control_net_lllite_patch: - def __init__(self, modules): + def __init__(self, modules: dict[str, nn.Module]): self.modules = modules def __call__(self, q, k, v, extra_options): module_pfx = extra_options_to_module_prefix(extra_options) - is_attn1 = q.shape[-1] == k.shape[-1] # self attention + is_attn1 = q.shape[-1] == k.shape[-1] if is_attn1: module_pfx = module_pfx + "_attn1" else: @@ -113,7 +111,7 @@ def load_control_net_lllite_patch(ctrl_sd, cond_image, multiplier, num_steps, st return control_net_lllite_patch(modules) -class LLLiteModule(torch.nn.Module): +class LLLiteModule(nn.Module): def __init__( self, name: str, @@ -126,7 +124,6 @@ class LLLiteModule(torch.nn.Module): num_steps: int, start_step: int, end_step: int, - dtype: torch.dtype = torch.float32 ): super().__init__() self.name = name @@ -138,60 +135,57 @@ class LLLiteModule(torch.nn.Module): self.is_first = False modules = [] - modules.append(torch.nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0)) # to latent (from VAE) size*2 + modules.append(nn.Conv2d(3, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0)) if depth == 1: - modules.append(torch.nn.ReLU(inplace=True)) - modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0)) + modules.append(nn.ReLU(inplace=True)) + modules.append(nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0)) elif depth == 2: - modules.append(torch.nn.ReLU(inplace=True)) - modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0)) + modules.append(nn.ReLU(inplace=True)) + modules.append(nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=4, stride=4, padding=0)) elif depth == 3: - # kernel size 8は大きすぎるので、4にする / kernel size 8 is too large, so set it to 4 - modules.append(torch.nn.ReLU(inplace=True)) - modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0)) - modules.append(torch.nn.ReLU(inplace=True)) - modules.append(torch.nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0)) + modules.append(nn.ReLU(inplace=True)) + modules.append(nn.Conv2d(cond_emb_dim // 2, cond_emb_dim // 2, kernel_size=4, stride=4, padding=0)) + modules.append(nn.ReLU(inplace=True)) + modules.append(nn.Conv2d(cond_emb_dim // 2, cond_emb_dim, kernel_size=2, stride=2, padding=0)) - self.conditioning1 = torch.nn.Sequential(*modules) + self.conditioning1 = nn.Sequential(*modules) if self.is_conv2d: - self.down = torch.nn.Sequential( - torch.nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0), - torch.nn.ReLU(inplace=True), + self.down = nn.Sequential( + nn.Conv2d(in_dim, mlp_dim, kernel_size=1, stride=1, padding=0), + nn.ReLU(inplace=True), ) - self.mid = torch.nn.Sequential( - torch.nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0), - torch.nn.ReLU(inplace=True), + self.mid = nn.Sequential( + nn.Conv2d(mlp_dim + cond_emb_dim, mlp_dim, kernel_size=1, stride=1, padding=0), + nn.ReLU(inplace=True), ) - self.up = torch.nn.Sequential( - torch.nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0), + self.up = nn.Sequential( + nn.Conv2d(mlp_dim, in_dim, kernel_size=1, stride=1, padding=0), ) else: - self.down = torch.nn.Sequential( - torch.nn.Linear(in_dim, mlp_dim), - torch.nn.ReLU(inplace=True), + self.down = nn.Sequential( + nn.Linear(in_dim, mlp_dim), + nn.ReLU(inplace=True), ) - self.mid = torch.nn.Sequential( - torch.nn.Linear(mlp_dim + cond_emb_dim, mlp_dim), - torch.nn.ReLU(inplace=True), + self.mid = nn.Sequential( + nn.Linear(mlp_dim + cond_emb_dim, mlp_dim), + nn.ReLU(inplace=True), ) - self.up = torch.nn.Sequential( - torch.nn.Linear(mlp_dim, in_dim), + self.up = nn.Sequential( + nn.Linear(mlp_dim, in_dim), ) self.depth = depth self.cond_image = None self.cond_emb = None self.current_step = 0 - self.to(dtype=dtype) - # @torch.inference_mode() def set_cond_image(self, cond_image): - # print("set_cond_image", self.name) self.cond_image = cond_image self.cond_emb = None self.current_step = 0 + @torch.inference_mode() def forward(self, x): if self.num_steps > 0: if self.current_step < self.start_step: @@ -199,36 +193,31 @@ class LLLiteModule(torch.nn.Module): return torch.zeros_like(x) elif self.current_step >= self.end_step: if self.is_first and self.current_step == self.end_step: - print(f"end LLLite: step {self.current_step}") + logger.debug(f"LLLite End: step {self.current_step}") self.current_step += 1 if self.current_step >= self.num_steps: - self.current_step = 0 # reset + self.current_step = 0 return torch.zeros_like(x) else: if self.is_first and self.current_step == self.start_step: - print(f"start LLLite: step {self.current_step}") + logger.debug(f"LLLite Start: step {self.current_step}") self.current_step += 1 if self.current_step >= self.num_steps: - self.current_step = 0 # reset + self.current_step = 0 if self.cond_emb is None: - # print(f"cond_emb is None, {self.name}") cx = self.conditioning1(self.cond_image.to(x.device, dtype=x.dtype)) if not self.is_conv2d: - # reshape / b,c,h,w -> b,h*w,c n, c, h, w = cx.shape cx = cx.view(n, c, h * w).permute(0, 2, 1) self.cond_emb = cx cx = self.cond_emb - # print(f"forward {self.name}, {cx.shape}, {x.shape}") - # uncond/condでxはバッチサイズが2倍 if x.shape[0] != cx.shape[0]: if self.is_conv2d: cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1, 1) else: - # print("x.shape[0] != cx.shape[0]", x.shape[0], cx.shape[0]) cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1) cx = torch.cat([cx, self.down(x)], dim=1 if self.is_conv2d else 2) @@ -238,35 +227,15 @@ class LLLiteModule(torch.nn.Module): class LLLiteLoader: - def __init__(self): - pass - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "model": ("MODEL",), - "model_name": None, - "cond_image": ("IMAGE",), - "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), - "steps": ("INT", {"default": 0, "min": 0, "max": 200, "step": 1}), - "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), - "end_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), - } - } + @staticmethod + def load_lllite(model: "UnetPatcher", state_dict, cond_image, strength, steps, start_percent, end_percent): + m = model.clone() - RETURN_TYPES = ("MODEL",) - FUNCTION = "load_lllite" - CATEGORY = "loaders" + patch = load_control_net_lllite_patch(state_dict, cond_image, strength, steps, start_percent, end_percent, model_dtype=model.model.diffusion_model.computation_dtype) - def load_lllite(self, model, state_dict, cond_image, strength, steps, start_percent, end_percent): - # cond_image is b,h,w,3, 0-1 - - model_lllite = model.clone() - patch = load_control_net_lllite_patch(state_dict, cond_image, strength, steps, start_percent, end_percent, - model.model.diffusion_model.computation_dtype) if patch is not None: - model_lllite.set_model_attn1_patch(patch) - model_lllite.set_model_attn2_patch(patch) + m.set_model_attn1_patch(patch) + m.set_model_attn2_patch(patch) - return (model_lllite,) + return m diff --git a/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite_anima.py b/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite_anima.py new file mode 100644 index 00000000..ba7e57e1 --- /dev/null +++ b/extensions-builtin/sd_forge_controlllite/lib_controllllite/lib_controllllite_anima.py @@ -0,0 +1,445 @@ +# https://github.com/kohya-ss/ComfyUI-Anima-LLLite/blob/main/control_net_lllite_anima.py + +import logging +import math +from typing import Final, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from backend.state_dict import load_state_dict + +logger = logging.getLogger("ControlNet") + + +# region Consts + + +TARGET_ATTENTION_CLASS: Final[str] = "SelfCrossAttention" +TARGET_MLP_CLASS: Final[str] = "GPT2FeedForward" + +ATOMIC_SPECIFIERS: Final[tuple[str]] = ( + "self_attn_q_pre", + "self_attn_kv_pre", + "cross_attn_q_pre", + "mlp_fc1_pre", +) + +PRESETS: Final[dict[str, tuple[str]]] = { + "self_attn_q": ("self_attn_q_pre",), + "self_attn_qkv": ("self_attn_q_pre", "self_attn_kv_pre"), + "self_attn_qkv_cross_q": ("self_attn_q_pre", "self_attn_kv_pre", "cross_attn_q_pre"), +} + +ASPP_DEFAULT_DILATIONS: Final[tuple[int]] = (1, 2, 4, 8) + + +_INTERNAL_MODULES_PREFIX = "lllite_modules." +_INTERNAL_COND_PREFIX = "conditioning1." +_INTERNAL_DEPTH_KEY = "depth_embeds" +_SAVED_COND_PREFIX = "lllite_conditioning1." +_SAVED_DEPTH_SUFFIX = ".depth_embed" + + +def parse_target_layers(spec: str) -> tuple[str]: + spec = spec.strip() + if spec in PRESETS: + return PRESETS[spec] + + parts = [p.strip() for p in spec.split(",") if p.strip()] + assert not any([p not in ATOMIC_SPECIFIERS for p in parts]) + + return tuple(a for a in ATOMIC_SPECIFIERS if a in parts) + + +# region Conditioning + + +def _gn(channels: int) -> nn.GroupNorm: + g = 8 + while g > 1 and channels % g != 0: + g //= 2 + return nn.GroupNorm(g, channels) + + +class _ResBlock(nn.Module): + def __init__(self, ch: int): + super().__init__() + self.norm1 = _gn(ch) + self.conv1 = nn.Conv2d(ch, ch, kernel_size=3, padding=1) + self.norm2 = _gn(ch) + self.conv2 = nn.Conv2d(ch, ch, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.conv1(F.silu(self.norm1(x))) + h = self.conv2(F.silu(self.norm2(h))) + return x + h + + +class _ASPP(nn.Module): + def __init__(self, ch: int, dilations: tuple[int] = ASPP_DEFAULT_DILATIONS): + super().__init__() + branches = [] + for d in dilations: + conv = nn.Conv2d(ch, ch, kernel_size=1) if d == 1 else nn.Conv2d(ch, ch, kernel_size=3, padding=d, dilation=d) + branches.append(nn.Sequential(conv, _gn(ch), nn.SiLU())) + self.branches = nn.ModuleList(branches) + self.global_pool = nn.AdaptiveAvgPool2d(1) + self.global_conv = nn.Sequential(nn.Conv2d(ch, ch, kernel_size=1), _gn(ch), nn.SiLU()) + n = len(dilations) + 1 + self.proj = nn.Sequential(nn.Conv2d(ch * n, ch, kernel_size=1), _gn(ch), nn.SiLU()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h, w = x.shape[-2:] + outs = [b(x) for b in self.branches] + g = self.global_conv(self.global_pool(x)) + g = F.interpolate(g, size=(h, w), mode="bilinear", align_corners=False) + outs.append(g) + return self.proj(torch.cat(outs, dim=1)) + + +class _Conditioning(nn.Module): + def __init__(self, cond_dim: int, cond_emb_dim: int, n_resblocks: int, use_aspp: bool = False, aspp_dilations: tuple[int] = ASPP_DEFAULT_DILATIONS): + super().__init__() + ch_half = cond_dim // 2 + self.conv1 = nn.Conv2d(3, ch_half, kernel_size=4, stride=4, padding=0) + self.norm1 = _gn(ch_half) + self.conv2 = nn.Conv2d(ch_half, ch_half, kernel_size=3, stride=1, padding=1) + self.norm2 = _gn(ch_half) + self.conv3 = nn.Conv2d(ch_half, cond_dim, kernel_size=4, stride=4, padding=0) + self.norm3 = _gn(cond_dim) + self.resblocks = nn.ModuleList([_ResBlock(cond_dim) for _ in range(n_resblocks)]) + self.aspp = _ASPP(cond_dim, aspp_dilations) if use_aspp else None + self.proj = nn.Conv2d(cond_dim, cond_emb_dim, kernel_size=1) + self.out_norm = nn.LayerNorm(cond_emb_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = F.silu(self.norm1(self.conv1(x))) + h = F.silu(self.norm2(self.conv2(h))) + h = F.silu(self.norm3(self.conv3(h))) + for rb in self.resblocks: + h = rb(h) + if self.aspp is not None: + h = self.aspp(h) + h = self.proj(h) + b, c, hh, ww = h.shape + h = h.view(b, c, hh * ww).permute(0, 2, 1).contiguous() + return self.out_norm(h) + + +# region Per-Linear LLLite Module + + +class LLLiteModuleDiT(nn.Module): + def __init__(self, name: str, org_module: nn.Linear, cond_emb_dim: int, mlp_dim: int, dropout: Optional[float] = None, multiplier: float = 1.0): + super().__init__() + self.lllite_name = name + + self.org_module: list[nn.Linear] = [org_module] + self.multiplier = multiplier + self.dropout = dropout + + in_dim = org_module.in_features + self.down = nn.Linear(in_dim, mlp_dim) + self.mid = nn.Linear(mlp_dim + cond_emb_dim, mlp_dim) + + self.cond_to_film = nn.Linear(cond_emb_dim, 2 * mlp_dim) + nn.init.zeros_(self.cond_to_film.weight) + nn.init.zeros_(self.cond_to_film.bias) + + self.up = nn.Linear(mlp_dim, in_dim) + nn.init.zeros_(self.up.weight) + nn.init.zeros_(self.up.bias) + + self.cond_emb: Optional[torch.Tensor] = None + self.org_forward = None + self.layer_idx: int = -1 + self._depth_embeds_ref: list[nn.Parameter] = [] + + self.num_steps: int = 0 + self.start_step: int = 0 + self.end_step: int = 0 + self.current_step: int = 0 + self.is_first: bool = False + + def apply_to(self): + if self.org_forward is None: + self.org_forward = self.org_module[0].forward + self.org_module[0].forward = self.forward + + def restore(self): + if self.org_forward is not None: + self.org_module[0].forward = self.org_forward + self.org_forward = None + + @torch.inference_mode() + def forward(self, x: torch.Tensor) -> torch.Tensor: + orig_shape = x.shape + is_5d = x.dim() == 5 + + def _pass(): + return self.org_forward(x.reshape(orig_shape) if is_5d else x) + + if self.multiplier == 0.0 or self.cond_emb is None: + return _pass() + + if self.num_steps > 0: + step = self.current_step + self.current_step += 1 + if self.current_step >= self.num_steps: + self.current_step = 0 + if step < self.start_step or step >= self.end_step: + return _pass() + + if is_5d: + B, T, H, W, D = orig_shape + x = x.reshape(B, T * H * W, D) + + cx = self.cond_emb + + if x.shape[0] != cx.shape[0]: + if x.shape[0] % cx.shape[0] != 0: + return self.org_forward(x.reshape(orig_shape) if is_5d else x) + cx = cx.repeat(x.shape[0] // cx.shape[0], 1, 1) + + if x.shape[1] != cx.shape[1]: + return self.org_forward(x.reshape(orig_shape) if is_5d else x) + + param_dtype = self.down.weight.dtype + x_proc = x if x.dtype == param_dtype else x.to(param_dtype) + if cx.dtype != param_dtype or cx.device != x.device: + cx = cx.to(device=x.device, dtype=param_dtype) + + if self._depth_embeds_ref: + depth_e = self._depth_embeds_ref[0][self.layer_idx] + if depth_e.dtype != param_dtype or depth_e.device != x.device: + depth_e = depth_e.to(device=x.device, dtype=param_dtype) + cond_local = cx + depth_e + else: + cond_local = cx + + h = F.silu(self.down(x_proc)) + gamma, beta = self.cond_to_film(cond_local).chunk(2, dim=-1) + m = self.mid(torch.cat([cond_local, h], dim=-1)) + m = F.silu(m * (1 + gamma) + beta) + if self.dropout is not None and self.training: + m = F.dropout(m, p=self.dropout) + + out = self.up(m) * self.multiplier + + if out.dtype != x.dtype: + out = out.to(x.dtype) + + y = self.org_forward(x + out) + if is_5d: + y = y.reshape(orig_shape[0], orig_shape[1], orig_shape[2], orig_shape[3], -1) + + return y + + +# region ControlNetLLLiteDiT + + +class ControlNetLLLiteDiT(nn.Module): + def __init__(self, dit: nn.Module, cond_emb_dim: int = 32, mlp_dim: int = 64, target_layers: str = "self_attn_q", dropout: Optional[float] = None, multiplier: float = 1.0, cond_dim: int = 64, cond_resblocks: int = 1, use_aspp: bool = False, aspp_dilations: tuple[int] = ASPP_DEFAULT_DILATIONS): + super().__init__() + atomics = parse_target_layers(target_layers) + self.multiplier = multiplier + self.target_atomics = atomics + + self.conditioning1 = _Conditioning( + cond_dim, + cond_emb_dim, + cond_resblocks, + use_aspp=use_aspp, + aspp_dilations=aspp_dilations, + ) + modules = self._create_modules(dit, cond_emb_dim, mlp_dim, atomics, dropout, multiplier) + self.lllite_modules: list[LLLiteModuleDiT] = nn.ModuleList(modules) + + n = len(self.lllite_modules) + self.depth_embeds = nn.Parameter(torch.zeros(n, cond_emb_dim)) + for i, m in enumerate(self.lllite_modules): + m.layer_idx = i + m._depth_embeds_ref = [self.depth_embeds] + + logger.info(f"Loaded Control-LLLite (Anima) ({n} modules)") + + @staticmethod + def _attn_atomic_match(is_self_attn: bool, child_name: str, atomics: tuple[str]) -> bool: + if "output_proj" in child_name: + return False + if is_self_attn: + if child_name == "q_proj": + return "self_attn_q_pre" in atomics + if child_name in ("k_proj", "v_proj"): + return "self_attn_kv_pre" in atomics + else: + if child_name == "q_proj": + return "cross_attn_q_pre" in atomics + return False + + def _create_modules(self, dit, cond_emb_dim, mlp_dim, atomics, dropout, multiplier): + modules = [] + want_mlp = "mlp_fc1_pre" in atomics + any_attn = any(a in atomics for a in ("self_attn_q_pre", "self_attn_kv_pre", "cross_attn_q_pre")) + + for name, module in dit.named_modules(): + cls = module.__class__.__name__ + + if any_attn and cls == TARGET_ATTENTION_CLASS: + if not hasattr(module, "is_SelfAttn"): + continue + is_self_attn = bool(module.is_SelfAttn) + for child_name, child in module.named_children(): + if not isinstance(child, nn.Linear): + continue + if not self._attn_atomic_match(is_self_attn, child_name, atomics): + continue + full_name = f"lllite_dit.{name}.{child_name}".replace(".", "_") + modules.append(LLLiteModuleDiT(full_name, child, cond_emb_dim, mlp_dim, dropout, multiplier)) + + elif want_mlp and cls == TARGET_MLP_CLASS: + child = getattr(module, "layer1", None) + if not isinstance(child, nn.Linear): + continue + full_name = f"lllite_dit.{name}.layer1".replace(".", "_") + modules.append(LLLiteModuleDiT(full_name, child, cond_emb_dim, mlp_dim, dropout, multiplier)) + + return modules + + def set_cond_image(self, cond_image: Optional[torch.Tensor]): + """cond_image: (B, 3, H, W) in [-1, 1]. None clears.""" + if cond_image is None: + for m in self.lllite_modules: + m.cond_emb = None + return + cx = self.conditioning1(cond_image) + for m in self.lllite_modules: + m.cond_emb = cx + + def set_multiplier(self, multiplier: float): + self.multiplier = multiplier + for m in self.lllite_modules: + m.multiplier = multiplier + + def set_step_range(self, num_steps: int, start_percent: float, end_percent: float): + start_step = math.floor(num_steps * start_percent) if start_percent > 0 else 0 + end_step = math.floor(num_steps * end_percent) if end_percent > 0 else num_steps + + for i, m in enumerate(self.lllite_modules): + m.num_steps = num_steps + m.start_step = start_step + m.end_step = end_step + m.current_step = 0 + m.is_first = i == 0 + + def apply_to(self): + for m in self.lllite_modules: + m.apply_to() + + def restore(self): + for m in self.lllite_modules: + m.restore() + self.set_cond_image(None) + + +# region Weight Loading (v2) + + +def _from_saved_state_dict(lllite: ControlNetLLLiteDiT, weights_sd: dict[str, torch.Tensor]) -> dict: + name_to_idx = {m.lllite_name: i for i, m in enumerate(lllite.lllite_modules)} + n_modules = len(name_to_idx) + out: dict = {} + depth_slices: dict = {} + + for k, v in weights_sd.items(): + if k.startswith(_SAVED_COND_PREFIX): + out[_INTERNAL_COND_PREFIX + k[len(_SAVED_COND_PREFIX) :]] = v + continue + if k.endswith(_SAVED_DEPTH_SUFFIX): + name = k[: -len(_SAVED_DEPTH_SUFFIX)] + if name in name_to_idx: + depth_slices[name_to_idx[name]] = v + continue + head, dot, tail = k.partition(".") + if dot and head in name_to_idx: + out[f"{_INTERNAL_MODULES_PREFIX}{name_to_idx[head]}.{tail}"] = v + continue + out[k] = v + + if depth_slices: + missing = [i for i in range(n_modules) if i not in depth_slices] + if missing: + raise RuntimeError(f"depth_embed slices missing for module indices: {missing}") + out[_INTERNAL_DEPTH_KEY] = torch.stack([depth_slices[i] for i in range(n_modules)], dim=0) + return out + + +def load_lllite_weights_from_dict(lllite: ControlNetLLLiteDiT, state_dict: dict[str, torch.Tensor]): + assert not any(k.startswith(_INTERNAL_MODULES_PREFIX) for k in state_dict) + converted = _from_saved_state_dict(lllite, state_dict) + load_state_dict(lllite, converted) + + +def infer_anima_config(state_dict: dict[str, torch.Tensor]) -> dict: + """Reconstruct ControlNetLLLiteDiT constructor kwargs from a saved state dict.""" + cond_emb_dim = 32 + cond_dim = 64 + mlp_dim = 64 + cond_resblocks = 0 + use_aspp = False + aspp_dilations = ASPP_DEFAULT_DILATIONS + + for k, v in state_dict.items(): + if k == "lllite_conditioning1.proj.weight": + cond_emb_dim = v.shape[0] + elif k == "lllite_conditioning1.conv1.weight": + cond_dim = v.shape[0] * 2 + + for k, v in state_dict.items(): + if k.endswith(".down.weight") and "conditioning1" not in k: + mlp_dim = v.shape[0] + break + + rb_indices = set() + for k in state_dict: + if "lllite_conditioning1.resblocks." in k: + parts = k.split(".") + try: + idx = parts.index("resblocks") + rb_indices.add(int(parts[idx + 1])) + except (ValueError, IndexError): + pass + + if rb_indices: + cond_resblocks = max(rb_indices) + 1 + + use_aspp = any("lllite_conditioning1.aspp" in k for k in state_dict) + + has_self_q = any("self_attn_q_proj.down.weight" in k for k in state_dict) + has_self_kv = any(("self_attn_k_proj" in k or "self_attn_v_proj" in k) and k.endswith(".down.weight") for k in state_dict) + has_cross_q = any("cross_attn_q_proj.down.weight" in k for k in state_dict) + has_mlp = any("_mlp_layer1.down.weight" in k for k in state_dict) + + parts = [] + if has_self_q: + parts.append("self_attn_q_pre") + if has_self_kv: + parts.append("self_attn_kv_pre") + if has_cross_q: + parts.append("cross_attn_q_pre") + if has_mlp: + parts.append("mlp_fc1_pre") + target_layers = ",".join(parts) if parts else "self_attn_q" + + return dict( + cond_emb_dim=cond_emb_dim, + mlp_dim=mlp_dim, + target_layers=target_layers, + cond_dim=cond_dim, + cond_resblocks=cond_resblocks, + use_aspp=use_aspp, + aspp_dilations=aspp_dilations, + ) diff --git a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py index 4a83a915..8484781e 100644 --- a/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py +++ b/extensions-builtin/sd_forge_controlllite/scripts/forge_controllllite.py @@ -1,38 +1,67 @@ +from lib_controllllite.lib_controllllite import LLLiteLoader +from lib_controllllite.lib_controllllite_anima import ( + ControlNetLLLiteDiT, + infer_anima_config, + load_lllite_weights_from_dict, +) + from modules_forge.shared import add_supported_control_model from modules_forge.supported_controlnet import ControlModelPatcher -from lib_controllllite.lib_controllllite import LLLiteLoader -opLLLiteLoader = LLLiteLoader().load_lllite +class ControlLLLiteAnimaPatcher(ControlModelPatcher): + @staticmethod + def try_build_from_state_dict(state_dict, ckpt_path): + if not any(k.startswith("lllite_dit") for k in state_dict): + return None + return ControlLLLiteAnimaPatcher(state_dict) + + def __init__(self, state_dict): + super().__init__() + self.state_dict = state_dict + self._lllite_net = None + + 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 + + if self._lllite_net is None: + dit = unet.model.diffusion_model + cfg = infer_anima_config(self.state_dict) + 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) + + cond_image = cond * 2.0 - 1.0 + 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) + self._lllite_net.apply_to() + + def process_after_every_sampling(self, *args, **kwargs): + if self._lllite_net is not None: + self._lllite_net.restore() class ControlLLLitePatcher(ControlModelPatcher): @staticmethod def try_build_from_state_dict(state_dict, ckpt_path): - if not any('lllite' in k for k in state_dict.keys()): + if not any(k.startswith("lllite") for k in state_dict): return None return ControlLLLitePatcher(state_dict) def __init__(self, state_dict): super().__init__() self.state_dict = state_dict - return def process_before_every_sampling(self, process, cond, mask, *args, **kwargs): unet = process.sd_model.forge_objects.unet - unet = opLLLiteLoader( - 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 - )[0] + 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 - return +add_supported_control_model(ControlLLLiteAnimaPatcher) + add_supported_control_model(ControlLLLitePatcher) diff --git a/modules_forge/supported_controlnet.py b/modules_forge/supported_controlnet.py index c3d4f6de..4dd121f0 100644 --- a/modules_forge/supported_controlnet.py +++ b/modules_forge/supported_controlnet.py @@ -77,7 +77,7 @@ class ControlNetPatcher(ControlModelPatcher): super().__init__(model_patcher) @staticmethod - def try_build_from_state_dict(controlnet_data, ckpt_path): + def try_build_from_state_dict(controlnet_data: dict[str, torch.Tensor], ckpt_path): if "lora_controlnet" in controlnet_data: return ControlNetPatcher(ControlLora(controlnet_data)) @@ -146,7 +146,8 @@ class ControlNetPatcher(ControlModelPatcher): else: net = load_t2i_adapter(controlnet_data) if net is None: - logger.error("Could not detect Control model type...") + if not any(k.startswith("lllite") for k in controlnet_data): # LLLite + logger.error("Could not detect Control model type...") return None return ControlNetPatcher(net)