mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
ControlNet LLLite
This commit is contained in:
Haoming
GitHub
parent
0db0d4d51f
commit
612fe1c8a8
@ -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
|
||||
|
||||
@ -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,
|
||||
)
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user