From ada263f2f6c9536687c5ddbe07fe4e9f4265f072 Mon Sep 17 00:00:00 2001 From: xiangyuanmeng Date: Wed, 24 Jul 2024 11:34:25 +0800 Subject: [PATCH] # Fix --- DeepCache_Fix.py | 249 +++++++++++++++++++++++++++++++++++++++++++++++ README.md | 32 ++++++ __init__.py | 7 ++ 3 files changed, 288 insertions(+) create mode 100755 DeepCache_Fix.py create mode 100644 README.md create mode 100755 __init__.py diff --git a/DeepCache_Fix.py b/DeepCache_Fix.py new file mode 100755 index 0000000..f9989b7 --- /dev/null +++ b/DeepCache_Fix.py @@ -0,0 +1,249 @@ +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"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..204ece7 --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# ComfyUI-DeepCache-Fix + +原始代码参考: https://gist.github.com/laksjdjf/435c512bc19636e9c9af4ee7bea9eb86 + +感谢 laksjdjf 分享的代码. + +将该库放在 ComfyUI/custom_nodes/ 下即可。 + +### 插件参数 + +- cache_interval: 缓存间隔, 单位: 步, 默认:3 +- cache_depth: 缓存深度, 默认:3 +- start_steps: 使用缓存的开始步数, 默认:0 +- end_steps: 使用缓存的结束步数, 默认:12 +- input_cache: 使用输入层缓存，默认:True 开启 +- middle_cahce: 使用中间层缓存，默认:True 开启e +- output_cache: 使用输出层缓存，默认:True 开启 + +### 举例(最佳实践) + +#### 模型 + +https://www.liblib.art/modelinfo/386109978c19484298d810d6f2830780 + +#### 生成 + +在总共15步的执行过程中，我们计划采取分阶段的策略。具体来说，如下 + +1. 前12步(start_steps=0, end_steps=12)将利用特定的插件来执行，以提高效率和效果。 +2. 从第13步开始，我们将切换回原始模型，完成剩下的3步。 + +这样的安排旨在结合两者的优势，确保整个流程的顺利进行。 \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100755 index 0000000..6746f0f --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +try: + import comfy.utils +except ImportError: + pass +else: + from .DeepCache_Fix import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']