mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
225 lines
8.6 KiB
Python
225 lines
8.6 KiB
Python
import os
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
|
|
from backend import memory_management
|
|
from backend.nn.cnets import cldm
|
|
from backend.operations import using_forge_operations
|
|
from backend.patcher.controlnet import (
|
|
ControlLora,
|
|
ControlNet,
|
|
apply_controlnet_advanced,
|
|
load_t2i_adapter,
|
|
logger,
|
|
)
|
|
from modules import shared
|
|
from modules_forge.packages.comfy.utils import unet_to_diffusers
|
|
from modules_forge.packages.huggingface_guess.detection import (
|
|
model_config_from_unet,
|
|
unet_config_from_diffusers_unet,
|
|
)
|
|
from modules_forge.shared import add_supported_control_model
|
|
|
|
|
|
class ModelCache:
|
|
def __init__(self):
|
|
self.cache = OrderedDict()
|
|
|
|
def get(self, key):
|
|
if key in self.cache:
|
|
self.cache.move_to_end(key)
|
|
return self.cache[key]
|
|
return None
|
|
|
|
def put(self, key, model):
|
|
self.cache[key] = model
|
|
self.cache.move_to_end(key)
|
|
|
|
max_size = max(1, shared.opts.data.get("control_net_model_cache_size", 3))
|
|
while len(self.cache) > max_size:
|
|
self.cache.popitem(last=False)
|
|
|
|
|
|
_CONTROL_MODEL_CACHE = ModelCache()
|
|
|
|
|
|
class ControlModelPatcher:
|
|
|
|
@staticmethod
|
|
def try_build_from_state_dict(state_dict, ckpt_path):
|
|
return None
|
|
|
|
def __init__(self, model_patcher=None):
|
|
self.model_patcher = model_patcher
|
|
self.strength = 1.0
|
|
self.start_percent = 0.0
|
|
self.end_percent = 1.0
|
|
self.positive_advanced_weighting = None
|
|
self.negative_advanced_weighting = None
|
|
self.advanced_frame_weighting = None
|
|
self.advanced_sigma_weighting = None
|
|
self.advanced_mask_weighting = None
|
|
|
|
def process_after_running_preprocessors(self, process, params, *args, **kwargs):
|
|
return
|
|
|
|
def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):
|
|
return
|
|
|
|
def process_after_every_sampling(self, process, params, *args, **kwargs):
|
|
return
|
|
|
|
|
|
class ControlNetPatcher(ControlModelPatcher):
|
|
|
|
def __init__(self, model_patcher):
|
|
super().__init__(model_patcher)
|
|
|
|
@staticmethod
|
|
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))
|
|
|
|
controlnet_config = None
|
|
if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: # diffusers format
|
|
unet_dtype = memory_management.unet_dtype()
|
|
controlnet_config = unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
|
|
diffusers_keys = unet_to_diffusers(controlnet_config)
|
|
diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
|
|
diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
|
|
|
|
count = 0
|
|
loop = True
|
|
while loop:
|
|
suffix = [".weight", ".bias"]
|
|
for s in suffix:
|
|
k_in = "controlnet_down_blocks.{}{}".format(count, s)
|
|
k_out = "zero_convs.{}.0{}".format(count, s)
|
|
if k_in not in controlnet_data:
|
|
loop = False
|
|
break
|
|
diffusers_keys[k_in] = k_out
|
|
count += 1
|
|
|
|
count = 0
|
|
loop = True
|
|
while loop:
|
|
suffix = [".weight", ".bias"]
|
|
for s in suffix:
|
|
if count == 0:
|
|
k_in = "controlnet_cond_embedding.conv_in{}".format(s)
|
|
else:
|
|
k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
|
|
k_out = "input_hint_block.{}{}".format(count * 2, s)
|
|
if k_in not in controlnet_data:
|
|
k_in = "controlnet_cond_embedding.conv_out{}".format(s)
|
|
loop = False
|
|
diffusers_keys[k_in] = k_out
|
|
count += 1
|
|
|
|
new_sd = {}
|
|
for k in diffusers_keys:
|
|
if k in controlnet_data:
|
|
new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
|
|
|
|
if "control_add_embedding.linear_1.bias" in controlnet_data: # Union Controlnet
|
|
controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0]
|
|
for k in list(controlnet_data.keys()):
|
|
new_k = k.replace(".attn.in_proj_", ".attn.in_proj.")
|
|
new_sd[new_k] = controlnet_data.pop(k)
|
|
|
|
leftover_keys = controlnet_data.keys()
|
|
if len(leftover_keys) > 0:
|
|
logger.warning("Leftover Keys: {}".format(leftover_keys))
|
|
controlnet_data = new_sd
|
|
|
|
pth_key = "control_model.zero_convs.0.0.weight"
|
|
pth = False
|
|
key = "zero_convs.0.0.weight"
|
|
if pth_key in controlnet_data:
|
|
pth = True
|
|
key = pth_key
|
|
prefix = "control_model."
|
|
elif key in controlnet_data:
|
|
prefix = ""
|
|
else:
|
|
net = load_t2i_adapter(controlnet_data)
|
|
if net is None:
|
|
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)
|
|
|
|
if controlnet_config is None:
|
|
unet_dtype = memory_management.unet_dtype()
|
|
controlnet_config = model_config_from_unet(controlnet_data, prefix, True).unet_config
|
|
controlnet_config["dtype"] = unet_dtype
|
|
|
|
load_device = memory_management.get_torch_device()
|
|
computation_dtype = shared.sd_model.forge_objects.unet.model.computation_dtype
|
|
controlnet_config.pop("out_channels")
|
|
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
|
|
controlnet_config["hint_width"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[0]
|
|
|
|
global _CONTROL_MODEL_CACHE
|
|
control_model = _CONTROL_MODEL_CACHE.get(ckpt_path)
|
|
|
|
if control_model is not None:
|
|
logger.info("Reusing ControlNet Model...")
|
|
else:
|
|
logger.info("Creating ControlNet Model...")
|
|
with using_forge_operations(dtype=unet_dtype, manual_cast_enabled=computation_dtype != unet_dtype):
|
|
control_model = cldm.ControlNet(**controlnet_config).to(dtype=unet_dtype)
|
|
|
|
if pth:
|
|
if "difference" in controlnet_data:
|
|
logger.warning("Please use an official Control model for better performance...")
|
|
|
|
class WeightsLoader(torch.nn.Module):
|
|
pass
|
|
|
|
w = WeightsLoader()
|
|
w.control_model = control_model
|
|
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
|
|
else:
|
|
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
|
|
|
|
if len(missing) > 0:
|
|
logger.warning("Missing ControlNet Keys: {}".format(missing))
|
|
|
|
if len(unexpected) > 0:
|
|
logger.debug("Unexpected ControlNet Keys: {}".format(unexpected))
|
|
|
|
_CONTROL_MODEL_CACHE.put(ckpt_path, control_model)
|
|
|
|
global_average_pooling = False
|
|
filename = os.path.splitext(ckpt_path)[0]
|
|
if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"):
|
|
# TODO: smarter way of enabling global_average_pooling
|
|
global_average_pooling = True
|
|
|
|
control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=computation_dtype)
|
|
return ControlNetPatcher(control)
|
|
|
|
def process_before_every_sampling(self, process, cond, mask, *args, control_type=None, **kwargs):
|
|
unet = process.sd_model.forge_objects.unet
|
|
unet = apply_controlnet_advanced(
|
|
unet=unet,
|
|
controlnet=self.model_patcher,
|
|
image_bchw=cond,
|
|
strength=self.strength,
|
|
start_percent=self.start_percent,
|
|
end_percent=self.end_percent,
|
|
positive_advanced_weighting=self.positive_advanced_weighting,
|
|
negative_advanced_weighting=self.negative_advanced_weighting,
|
|
advanced_frame_weighting=self.advanced_frame_weighting,
|
|
advanced_sigma_weighting=self.advanced_sigma_weighting,
|
|
advanced_mask_weighting=self.advanced_mask_weighting,
|
|
control_type=control_type,
|
|
)
|
|
process.sd_model.forge_objects.unet = unet
|
|
|
|
|
|
add_supported_control_model(ControlNetPatcher)
|