IC-Light 在神经网络框架中的集成
导入语句
import torch
import torchvision.transforms as transforms
import folder_paths
import os
import types
import numpy as np
import torch.nn.functional as F
from comfy.utils import load_torch_file
from .utils.convert_unet import convert_iclight_unet
from .utils.patches import calculate_weight_adjust_channel
from .utils.image import generate_gradient_image, LightPosition
from nodes import MAX_RESOLUTION
from comfy.model_patcher import ModelPatcher
import model_management
这些导入语句提供了执行数组操作、深度学习模型操作、文件路径处理等功能所需的各种模块和函数。特别是 torch
和 torchvision
库用于深度学习相关的操作,而 folder_paths
和 os
用于文件系统操作。
类 LoadAndApplyICLightUnet
的定义
输入类型定义
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"model_path": (folder_paths.get_filename_list("unet"), )
}
}
此方法定义了必需的输入类型,包括一个模型对象和一个模型路径,后者使用 folder_paths.get_filename_list("unet")
函数获取UNet模型文件的列表。
加载函数
def load(self, model, model_path):
type_str = str(type(model.model.model_config).__name__)
if "SD15" not in type_str:
raise Exception(f"Attempted to load {type_str} model, IC-Light is only compatible with SD 1.5 models.")
print("LoadAndApplyICLightUnet: Checking IC-Light Unet path")
model_full_path = folder_paths.get_full_path("unet", model_path)
if not os.path.exists(model_full_path):
raise Exception("Invalid model path")
else:
print("LoadAndApplyICLightUnet: Loading IC-Light Unet weights")
model_clone = model.clone()
iclight_state_dict = load_torch_file(model_full_path)
print("LoadAndApplyICLightUnet: Attempting to add patches with IC-Light Unet weights")
try:
if 'conv_in.weight' in iclight_state_dict:
iclight_state_dict = convert_iclight_unet(iclight_state_dict)
in_channels = iclight_state_dict["diffusion_model.input_blocks.0.0.weight"].shape[1]
for key in iclight_state_dict:
model_clone.add_patches({key: (iclight_state_dict[key],)}, 1.0, 1.0)
else:
for key in iclight_state_dict:
model_clone.add_patches({"diffusion_model." + key: (iclight_state_dict[key],)}, 1.0, 1.0)
in_channels = iclight_state_dict["input_blocks.0.0.weight"].shape[1]
except:
raise Exception("Could not patch model")
print("LoadAndApplyICLightUnet: Added LoadICLightUnet patches")
try:
ModelPatcher.calculate_weight = calculate_weight_adjust_channel(ModelPatcher.calculate_weight)
except:
raise Exception("IC-Light: Could not patch calculate_weight")
# Mimic the existing IP2P class to enable extra_conds
def bound_extra_conds(self, **kwargs):
return ICLight.extra_conds(self, **kwargs)
new_extra_conds = types.MethodType(bound_extra_conds, model_clone.model)
model_clone.add_object_patch("extra_conds", new_extra_conds)
model_clone.model.model_config.unet_config["in_channels"] = in_channels
return (model_clone, )
此函数负责加载并应用IC-Light的UNet模型。它首先检查模型是否符合兼容性(仅限SD 1.5模型),然后加载模型权重,并尝试将这些权重作为补丁应用到克隆的模型上。此外,它还修改了模型的权重计算方法,并添加了额外的条件处理方法以增强功能。
类 ICLightConditioning
的定义
输入类型定义
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"positive": ("CONDITIONING", ),
"negative": ("CONDITIONING", ),
"vae": ("VAE", ),
"foreground": ("LATENT", ),
"multiplier": ("FLOAT", {"default": 0.18215, "min": 0.0, "max": 1.0, "step": 0.001}),
},
"optional": {
"opt_background": ("LATENT", ),
},
}
定义了IC-Light调节所需的输入参数,包括正负调节条件、VAE模型、前景潜像、乘数以及可选的背景潜像。
编码函数
def encode(self, positive, negative, vae, foreground, multiplier, opt_background=None):
samples_1 = foreground["samples"]
if opt_background is not None:
samples_2 = opt_background["samples"]
repeats_1 = samples_2.size(0) // samples_1.size(0)
repeats_2 = samples_1.size(0) // samples_2.size(0)
if samples_1.shape[1:] != samples_2.shape[1:]:
samples_2 = comfy.utils.common_upscale(samples_2, samples_1.shape[-1], samples_1.shape[-2], "bilinear", "disabled")
if repeats_1 > 1:
samples_1 = samples_1.repeat(repeats_1, 1, 1, 1)
if repeats_2 > 1:
samples_2 = samples_2.repeat(repeats_2, 1, 1, 1)
concat_latent = torch.cat((samples_1, samples_2), dim=1)
else:
concat_latent = samples_1
out_latent = torch.zeros_like(samples_1)
out = []
for conditioning in [positive, negative]:
c = []
for t in conditioning:
d = t[1].copy()
d["concat_latent_image"] = concat_latent * multiplier
n = [t[0], d]
c.append(n)
out.append(c)
return (out[0], out[1], {"samples": out_latent})
此函数处理传入的正负条件,应用VAE和潜像,并将乘数应用于组合潜像以生成新的调节结果。它支持处理可选的背景潜像,允许更复杂的图像合成操作。
节点类和显示名称的注册
NODE_CLASS_MAPPINGS = {
"LoadAndApplyICLightUnet": LoadAndApplyICLightUnet,
"ICLightConditioning": ICLightConditioning,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"LoadAndApplyICLightUnet": "Load And Apply IC-Light",
"ICLightConditioning": "IC-Light Conditioning",
}
这些代码行将上述定义的类注册为可用的节点类型,并设置它们的显示名称,使得这些功能可以在更大的系统或框架中通过指定的名称被调用。