321 lines
12 KiB
Python
321 lines
12 KiB
Python
from typing import List, NamedTuple, Any
|
|
import numpy as np
|
|
import cv2
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
from diffusers import UNet2DConditionModel
|
|
from diffusers.models.unet_2d_condition import UNet2DConditionOutput
|
|
|
|
import library.model_util as model_util
|
|
|
|
|
|
class ControlNetInfo(NamedTuple):
|
|
unet: Any
|
|
net: Any
|
|
prep: Any
|
|
weight: float
|
|
ratio: float
|
|
|
|
|
|
class ControlNet(torch.nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
# make control model
|
|
self.control_model = torch.nn.Module()
|
|
|
|
dims = [320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280]
|
|
zero_convs = torch.nn.ModuleList()
|
|
for i, dim in enumerate(dims):
|
|
sub_list = torch.nn.ModuleList([torch.nn.Conv2d(dim, dim, 1)])
|
|
zero_convs.append(sub_list)
|
|
self.control_model.add_module("zero_convs", zero_convs)
|
|
|
|
middle_block_out = torch.nn.Conv2d(1280, 1280, 1)
|
|
self.control_model.add_module("middle_block_out", torch.nn.ModuleList([middle_block_out]))
|
|
|
|
dims = [16, 16, 32, 32, 96, 96, 256, 320]
|
|
strides = [1, 1, 2, 1, 2, 1, 2, 1]
|
|
prev_dim = 3
|
|
input_hint_block = torch.nn.Sequential()
|
|
for i, (dim, stride) in enumerate(zip(dims, strides)):
|
|
input_hint_block.append(torch.nn.Conv2d(prev_dim, dim, 3, stride, 1))
|
|
if i < len(dims) - 1:
|
|
input_hint_block.append(torch.nn.SiLU())
|
|
prev_dim = dim
|
|
self.control_model.add_module("input_hint_block", input_hint_block)
|
|
|
|
|
|
def load_control_net(v2, unet, model):
|
|
device = unet.device
|
|
|
|
# control sdからキー変換しつつU-Netに対応する部分のみ取り出し、DiffusersのU-Netに読み込む
|
|
# state dictを読み込む
|
|
print(f"ControlNet: loading control SD model : {model}")
|
|
|
|
if model_util.is_safetensors(model):
|
|
ctrl_sd_sd = load_file(model)
|
|
else:
|
|
ctrl_sd_sd = torch.load(model, map_location='cpu')
|
|
ctrl_sd_sd = ctrl_sd_sd.pop("state_dict", ctrl_sd_sd)
|
|
|
|
# 重みをU-Netに読み込めるようにする。ControlNetはSD版のstate dictなので、それを読み込む
|
|
is_difference = "difference" in ctrl_sd_sd
|
|
print("ControlNet: loading difference")
|
|
|
|
# ControlNetには存在しないキーがあるので、まず現在のU-NetでSD版の全keyを作っておく
|
|
# またTransfer Controlの元weightとなる
|
|
ctrl_unet_sd_sd = model_util.convert_unet_state_dict_to_sd(v2, unet.state_dict())
|
|
|
|
# 元のU-Netに影響しないようにコピーする。またprefixが付いていないので付ける
|
|
for key in list(ctrl_unet_sd_sd.keys()):
|
|
ctrl_unet_sd_sd["model.diffusion_model." + key] = ctrl_unet_sd_sd.pop(key).clone()
|
|
|
|
zero_conv_sd = {}
|
|
for key in list(ctrl_sd_sd.keys()):
|
|
if key.startswith("control_"):
|
|
unet_key = "model.diffusion_" + key[len("control_"):]
|
|
if unet_key not in ctrl_unet_sd_sd: # zero conv
|
|
zero_conv_sd[key] = ctrl_sd_sd[key]
|
|
continue
|
|
if is_difference: # Transfer Control
|
|
ctrl_unet_sd_sd[unet_key] += ctrl_sd_sd[key].to(device, dtype=unet.dtype)
|
|
else:
|
|
ctrl_unet_sd_sd[unet_key] = ctrl_sd_sd[key].to(device, dtype=unet.dtype)
|
|
|
|
unet_config = model_util.create_unet_diffusers_config(v2)
|
|
ctrl_unet_du_sd = model_util.convert_ldm_unet_checkpoint(v2, ctrl_unet_sd_sd, unet_config) # DiffUsers版ControlNetのstate dict
|
|
|
|
# ControlNetのU-Netを作成する
|
|
ctrl_unet = UNet2DConditionModel(**unet_config)
|
|
info = ctrl_unet.load_state_dict(ctrl_unet_du_sd)
|
|
print("ControlNet: loading Control U-Net:", info)
|
|
|
|
# U-Net以外のControlNetを作成する
|
|
# TODO support middle only
|
|
ctrl_net = ControlNet()
|
|
info = ctrl_net.load_state_dict(zero_conv_sd)
|
|
print("ControlNet: loading ControlNet:", info)
|
|
|
|
ctrl_unet.to(unet.device, dtype=unet.dtype)
|
|
ctrl_net.to(unet.device, dtype=unet.dtype)
|
|
return ctrl_unet, ctrl_net
|
|
|
|
|
|
def load_preprocess(prep_type: str):
|
|
if prep_type is None or prep_type.lower() == "none":
|
|
return None
|
|
|
|
if prep_type.startswith("canny"):
|
|
args = prep_type.split("_")
|
|
th1 = int(args[1]) if len(args) >= 2 else 63
|
|
th2 = int(args[2]) if len(args) >= 3 else 191
|
|
|
|
def canny(img):
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
return cv2.Canny(img, th1, th2)
|
|
return canny
|
|
|
|
print("Unsupported prep type:", prep_type)
|
|
return None
|
|
|
|
|
|
def preprocess_ctrl_net_hint_image(image):
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = image[:, :, ::-1].copy() # rgb to bgr
|
|
image = image[None].transpose(0, 3, 1, 2) # nchw
|
|
image = torch.from_numpy(image)
|
|
return image # 0 to 1
|
|
|
|
|
|
def get_guided_hints(control_nets: List[ControlNetInfo], num_latent_input, b_size, hints):
|
|
guided_hints = []
|
|
for i, cnet_info in enumerate(control_nets):
|
|
# hintは 1枚目の画像のcnet1, 1枚目の画像のcnet2, 1枚目の画像のcnet3, 2枚目の画像のcnet1, 2枚目の画像のcnet2 ... と並んでいること
|
|
b_hints = []
|
|
if len(hints) == 1: # すべて同じ画像をhintとして使う
|
|
hint = hints[0]
|
|
if cnet_info.prep is not None:
|
|
hint = cnet_info.prep(hint)
|
|
hint = preprocess_ctrl_net_hint_image(hint)
|
|
b_hints = [hint for _ in range(b_size)]
|
|
else:
|
|
for bi in range(b_size):
|
|
hint = hints[(bi * len(control_nets) + i) % len(hints)]
|
|
if cnet_info.prep is not None:
|
|
hint = cnet_info.prep(hint)
|
|
hint = preprocess_ctrl_net_hint_image(hint)
|
|
b_hints.append(hint)
|
|
b_hints = torch.cat(b_hints, dim=0)
|
|
b_hints = b_hints.to(cnet_info.unet.device, dtype=cnet_info.unet.dtype)
|
|
|
|
guided_hint = cnet_info.net.control_model.input_hint_block(b_hints)
|
|
guided_hints.append(guided_hint)
|
|
return guided_hints
|
|
|
|
|
|
def call_unet_and_control_net(step, num_latent_input, original_unet, control_nets: List[ControlNetInfo], guided_hints, current_ratio, sample, timestep, encoder_hidden_states):
|
|
# ControlNet
|
|
# 複数のControlNetの場合は、出力をマージするのではなく交互に適用する
|
|
cnet_cnt = len(control_nets)
|
|
cnet_idx = step % cnet_cnt
|
|
cnet_info = control_nets[cnet_idx]
|
|
|
|
# print(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
|
|
if cnet_info.ratio < current_ratio:
|
|
return original_unet(sample, timestep, encoder_hidden_states)
|
|
|
|
guided_hint = guided_hints[cnet_idx]
|
|
guided_hint = guided_hint.repeat((num_latent_input, 1, 1, 1))
|
|
outs = unet_forward(True, cnet_info.net, cnet_info.unet, guided_hint, None, sample, timestep, encoder_hidden_states)
|
|
outs = [o * cnet_info.weight for o in outs]
|
|
|
|
# U-Net
|
|
return unet_forward(False, cnet_info.net, original_unet, None, outs, sample, timestep, encoder_hidden_states)
|
|
|
|
|
|
"""
|
|
# これはmergeのバージョン
|
|
# ControlNet
|
|
cnet_outs_list = []
|
|
for i, cnet_info in enumerate(control_nets):
|
|
# print(current_ratio, cnet_info.prep, cnet_info.weight, cnet_info.ratio)
|
|
if cnet_info.ratio < current_ratio:
|
|
continue
|
|
guided_hint = guided_hints[i]
|
|
outs = unet_forward(True, cnet_info.net, cnet_info.unet, guided_hint, None, sample, timestep, encoder_hidden_states)
|
|
for i in range(len(outs)):
|
|
outs[i] *= cnet_info.weight
|
|
|
|
cnet_outs_list.append(outs)
|
|
|
|
count = len(cnet_outs_list)
|
|
if count == 0:
|
|
return original_unet(sample, timestep, encoder_hidden_states)
|
|
|
|
# sum of controlnets
|
|
for i in range(1, count):
|
|
cnet_outs_list[0] += cnet_outs_list[i]
|
|
|
|
# U-Net
|
|
return unet_forward(False, cnet_info.net, original_unet, None, cnet_outs_list[0], sample, timestep, encoder_hidden_states)
|
|
"""
|
|
|
|
|
|
def unet_forward(is_control_net, control_net: ControlNet, unet: UNet2DConditionModel, guided_hint, ctrl_outs, sample, timestep, encoder_hidden_states):
|
|
# copy from UNet2DConditionModel
|
|
default_overall_up_factor = 2**unet.num_upsamplers
|
|
|
|
forward_upsample_size = False
|
|
upsample_size = None
|
|
|
|
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
|
print("Forward upsample size to force interpolation output size.")
|
|
forward_upsample_size = True
|
|
|
|
# 0. center input if necessary
|
|
if unet.config.center_input_sample:
|
|
sample = 2 * sample - 1.0
|
|
|
|
# 1. time
|
|
timesteps = timestep
|
|
if not torch.is_tensor(timesteps):
|
|
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
# This would be a good case for the `match` statement (Python 3.10+)
|
|
is_mps = sample.device.type == "mps"
|
|
if isinstance(timestep, float):
|
|
dtype = torch.float32 if is_mps else torch.float64
|
|
else:
|
|
dtype = torch.int32 if is_mps else torch.int64
|
|
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
|
elif len(timesteps.shape) == 0:
|
|
timesteps = timesteps[None].to(sample.device)
|
|
|
|
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
timesteps = timesteps.expand(sample.shape[0])
|
|
|
|
t_emb = unet.time_proj(timesteps)
|
|
|
|
# timesteps does not contain any weights and will always return f32 tensors
|
|
# but time_embedding might actually be running in fp16. so we need to cast here.
|
|
# there might be better ways to encapsulate this.
|
|
t_emb = t_emb.to(dtype=unet.dtype)
|
|
emb = unet.time_embedding(t_emb)
|
|
|
|
outs = [] # output of ControlNet
|
|
zc_idx = 0
|
|
|
|
# 2. pre-process
|
|
sample = unet.conv_in(sample)
|
|
if is_control_net:
|
|
sample += guided_hint
|
|
outs.append(control_net.control_model.zero_convs[zc_idx][0](sample)) # , emb, encoder_hidden_states))
|
|
zc_idx += 1
|
|
|
|
# 3. down
|
|
down_block_res_samples = (sample,)
|
|
for downsample_block in unet.down_blocks:
|
|
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
|
sample, res_samples = downsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
)
|
|
else:
|
|
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
|
if is_control_net:
|
|
for rs in res_samples:
|
|
outs.append(control_net.control_model.zero_convs[zc_idx][0](rs)) # , emb, encoder_hidden_states))
|
|
zc_idx += 1
|
|
|
|
down_block_res_samples += res_samples
|
|
|
|
# 4. mid
|
|
sample = unet.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
|
|
if is_control_net:
|
|
outs.append(control_net.control_model.middle_block_out[0](sample))
|
|
return outs
|
|
|
|
if not is_control_net:
|
|
sample += ctrl_outs.pop()
|
|
|
|
# 5. up
|
|
for i, upsample_block in enumerate(unet.up_blocks):
|
|
is_final_block = i == len(unet.up_blocks) - 1
|
|
|
|
res_samples = down_block_res_samples[-len(upsample_block.resnets):]
|
|
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
|
|
|
if not is_control_net and len(ctrl_outs) > 0:
|
|
res_samples = list(res_samples)
|
|
apply_ctrl_outs = ctrl_outs[-len(res_samples):]
|
|
ctrl_outs = ctrl_outs[:-len(res_samples)]
|
|
for j in range(len(res_samples)):
|
|
res_samples[j] = res_samples[j] + apply_ctrl_outs[j]
|
|
res_samples = tuple(res_samples)
|
|
|
|
# if we have not reached the final block and need to forward the
|
|
# upsample size, we do it here
|
|
if not is_final_block and forward_upsample_size:
|
|
upsample_size = down_block_res_samples[-1].shape[2:]
|
|
|
|
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
|
sample = upsample_block(
|
|
hidden_states=sample,
|
|
temb=emb,
|
|
res_hidden_states_tuple=res_samples,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
upsample_size=upsample_size,
|
|
)
|
|
else:
|
|
sample = upsample_block(
|
|
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
|
|
)
|
|
# 6. post-process
|
|
sample = unet.conv_norm_out(sample)
|
|
sample = unet.conv_act(sample)
|
|
sample = unet.conv_out(sample)
|
|
|
|
return UNet2DConditionOutput(sample=sample)
|