# Fix

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"]

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步。 
+
+这样的安排旨在结合两者的优势，确保整个流程的顺利进行。

__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']