250 lines
10 KiB
Python
Executable File
250 lines
10 KiB
Python
Executable File
import torch
|
|
from comfy.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, timestep_embedding, th, apply_control
|
|
|
|
class DeepCache_Fix:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
|
|
"""
|
|
静态方法，用于定义输入参数的类型和默认值。
|
|
|
|
该方法返回一个字典，其中包含了不同输入参数的配置信息。每个输入参数都是一个键值对，
|
|
键表示参数名，值是一个元组，包含参数的类型和一个字典，该字典描述了参数的更多细节，
|
|
如默认值、最小值、最大值等。
|
|
|
|
返回:
|
|
dict: 包含所有输入参数配置的字典。
|
|
"""
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"cache_interval": ("INT", {
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 1000,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"cache_depth": ("INT", {
|
|
"default": 3,
|
|
"min": 0,
|
|
"max": 12,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"start_steps": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"end_steps": ("INT", {
|
|
"default": 12,
|
|
"min": 0,
|
|
"max": 100,
|
|
"step": 1,
|
|
"display": "number"
|
|
}),
|
|
"input_cache": (["No", "Yes"], {"default":"Yes"}),
|
|
"middle_cahce": (["No", "Yes"], {"default":"Yes"}),
|
|
"output_cache": (["No", "Yes"], {"default":"Yes"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "loaders"
|
|
|
|
def apply(self, model, cache_interval, cache_depth, start_steps, end_steps, input_cache, middle_cahce, output_cache):
|
|
# 初始化一些变量
|
|
current_time = -1
|
|
current_step = -1
|
|
model_step = 0
|
|
cache_h = None
|
|
|
|
# 创建一个新的模型副本，用于存储修改后的模型。
|
|
new_model = model.clone()
|
|
# 获取并初始化模型的扩散部分。
|
|
unet = new_model.model.diffusion_model
|
|
# 获取并初始化模型的扩散部分。
|
|
dtype = new_model.model.get_dtype()
|
|
|
|
def cache_apply_methon(current, start, end):
|
|
"""
|
|
判断当前步骤是否在指定的开始步骤和结束步骤范围内。
|
|
|
|
参数:
|
|
current -- 当前的步骤数
|
|
start -- 范围的起始步骤数
|
|
end -- 范围的结束步骤数
|
|
|
|
返回:
|
|
如果当前步骤在指定范围内，则返回True；否则返回False。
|
|
"""
|
|
return start <= current <= end
|
|
|
|
def apply_model(model_function, kwargs):
|
|
"""
|
|
应用模型函数到给定的输入上。
|
|
|
|
这个函数处理模型的输入和输出，包括数据类型转换、条件的添加和模型的分块计算。
|
|
它还处理缓存机制，以在多次调用之间复用计算结果，提高效率。
|
|
|
|
:param model_function: 模型函数，一个接受kwargs参数的函数。
|
|
:param kwargs: 包含模型输入和配置的字典。包括输入数据、时间步、条件等。
|
|
:return: 模型处理后的输出。
|
|
"""
|
|
|
|
# 声明一些非局部变量，用于处理缓存和当前时间步等状态。
|
|
nonlocal model_step, cache_h, current_time, current_step
|
|
|
|
# 从kwargs中提取必要的输入和配置。
|
|
xa = kwargs["input"]
|
|
t = kwargs["timestep"]
|
|
c_concat = kwargs["c"].get("c_concat", None)
|
|
c_crossattn = kwargs["c"].get("c_crossattn", None)
|
|
y = kwargs["c"].get("y", None)
|
|
control = kwargs["c"].get("control", None)
|
|
transformer_options = kwargs["c"].get("transformer_options", None)
|
|
|
|
# 根据当前时间步计算输入xc。
|
|
sigma = t
|
|
xc = new_model.model.model_sampling.calculate_input(sigma, xa)
|
|
if c_concat is not None:
|
|
# 将输入xc与跨注意力的上下文c_concat进行拼接。
|
|
xc = torch.cat([xc] + [c_concat], dim=1)
|
|
|
|
# 处理跨注意力的上下文和数据类型的转换。
|
|
context = c_crossattn
|
|
xc = xc.to(dtype)
|
|
# 将时间步转换为指定的数据类型。
|
|
t = new_model.model.model_sampling.timestep(t).float()
|
|
context = context.to(dtype)
|
|
|
|
# 将所有额外的条件转换为指定的数据类型。
|
|
extra_conds = {}
|
|
for o in kwargs:
|
|
extra = kwargs[o]
|
|
if hasattr(extra, "to"):
|
|
extra = extra.to(dtype)
|
|
extra_conds[o] = extra
|
|
|
|
# 初始化模型的输入和配置。
|
|
x = xc
|
|
timesteps = t
|
|
y = None if y is None else y.to(dtype)
|
|
transformer_options["original_shape"] = list(x.shape)
|
|
transformer_options["current_index"] = 0
|
|
transformer_patches = transformer_options.get("patches", {})
|
|
|
|
model_step += 1
|
|
# 更新当前时间步和缓存状态，根据当前时间步决定是否应用模型。
|
|
if t[0].item() > current_time:
|
|
model_step = 0
|
|
current_step = -1
|
|
# 判断是否需要应用模型，根据当前时间步和指定的时间范围。
|
|
cache_apply = cache_apply_methon(model_step, start_steps, end_steps)
|
|
if cache_apply:
|
|
current_step += 1
|
|
else:
|
|
current_step = -1
|
|
current_time = t[0].item()
|
|
# print(f"model_step: {model_step}, {cache_apply}")
|
|
|
|
# 确保如果模型是分类的，那么必须提供标签y。
|
|
assert (y is not None) == (
|
|
unet.num_classes is not None
|
|
), "must specify y if and only if the model is class-conditional"
|
|
|
|
# 处理时间嵌入和模型的输入、中间和输出块。
|
|
hs = []
|
|
t_emb = timestep_embedding(timesteps, unet.model_channels, repeat_only=False).to(unet.dtype)
|
|
emb = unet.time_embed(t_emb)
|
|
if unet.num_classes is not None:
|
|
assert y.shape[0] == x.shape[0]
|
|
emb = emb + unet.label_emb(y)
|
|
xuh = x.type(unet.dtype)
|
|
# current_step 是 cache_interval 的整数倍?
|
|
step_cache_interval = current_step % cache_interval
|
|
# 循环处理输入块。
|
|
for id, module in enumerate(unet.input_blocks):
|
|
transformer_options["block"] = ("input", id)
|
|
xuh = forward_timestep_embed(module, xuh, emb, context, transformer_options)
|
|
xuh = apply_control(xuh, control, 'input')
|
|
if "input_block_patch" in transformer_patches:
|
|
patch = transformer_patches["input_block_patch"]
|
|
for p in patch:
|
|
xuh = p(xuh, transformer_options)
|
|
hs.append(xuh)
|
|
if "input_block_patch_after_skip" in transformer_patches:
|
|
patch = transformer_patches["input_block_patch_after_skip"]
|
|
for p in patch:
|
|
xuh = p(xuh, transformer_options)
|
|
|
|
# 根据缓存策略决定是否继续处理或使用缓存。
|
|
if id == cache_depth and cache_apply and input_cache:
|
|
if not step_cache_interval == 0:
|
|
break
|
|
|
|
# 处理中间块，同样考虑缓存策略。
|
|
# 如果 current_step 是 cache_interval 的整数倍
|
|
# 或者 cache_apply 为 False
|
|
# 或者 middle_cahce 为 False (开关关闭)
|
|
# 则执行中间块的处理。
|
|
if step_cache_interval == 0 or not cache_apply or not middle_cahce:
|
|
transformer_options["block"] = ("middle", 0)
|
|
xuh = forward_timestep_embed(unet.middle_block, xuh, emb, context, transformer_options)
|
|
xuh = apply_control(xuh, control, 'middle')
|
|
|
|
# 处理输出块，包括缓存的加载和使用。
|
|
for id, module in enumerate(unet.output_blocks):
|
|
if id < len(unet.output_blocks) - cache_depth - 1 and cache_apply and output_cache:
|
|
if not step_cache_interval == 0:
|
|
continue
|
|
if id == len(unet.output_blocks) - cache_depth - 1 and cache_apply and output_cache:
|
|
if step_cache_interval == 0:
|
|
cache_h = xuh # cache
|
|
else:
|
|
xuh = cache_h # load cache
|
|
transformer_options["block"] = ("output", id)
|
|
hsp = hs.pop()
|
|
hsp = apply_control(hsp, control, 'output')
|
|
if "output_block_patch" in transformer_patches:
|
|
patch = transformer_patches["output_block_patch"]
|
|
for p in patch:
|
|
xuh, hsp = p(xuh, hsp, transformer_options)
|
|
xuh = th.cat([xuh, hsp], dim=1)
|
|
del hsp
|
|
if len(hs) > 0:
|
|
output_shape = hs[-1].shape
|
|
else:
|
|
output_shape = None
|
|
xuh = forward_timestep_embed(module, xuh, emb, context, transformer_options, output_shape)
|
|
|
|
# 将输出转换回原始数据类型，并根据模型配置进行噪声消除计算。
|
|
xuh = xuh.type(x.dtype)
|
|
if unet.predict_codebook_ids:
|
|
model_output = unet.id_predictor(xuh)
|
|
else:
|
|
model_output = unet.out(xuh)
|
|
|
|
# 返回计算得到的最终输出。
|
|
return new_model.model.model_sampling.calculate_denoised(sigma, model_output, xa)
|
|
|
|
new_model.set_model_unet_function_wrapper(apply_model)
|
|
|
|
return (new_model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeepCache_Fix": DeepCache_Fix,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DeepCache_Fix": "DeepCache_Fix",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|