first commit

README.md

@@ -0,0 +1,10 @@
+# TeaCache技术现已整合到ComfyUI！🚀🚀🚀
+TeaCache是什么？——TeaCache是一种Diffusion Cache技术，它无需训练即可以几乎无损的方式大幅加速Diffusion Model推理🚀🚀🚀
+# ComfyUI-TeaCache重要更新：
+- TeaCache现已整合到ComfyUI，与官方原生Diffusion节点兼容！
+- TeaCache使用简单，仅需把我实现的TeaCache节点与comfyui官方节点相连即可无缝使用！
+- 经测试，对于FLUX模型，TeaCache可实现1.4倍无损加速，2倍加速无明显质量损失！
+- 支持LoRA！
+- 支持ControlNet!
+
+Enjoy It!!!

__init__.py

@@ -0,0 +1,10 @@
+# only import if running as a custom node
+try:
+    import comfy.utils
+except ImportError:
+    pass
+else:
+    from .nodes import NODE_CLASS_MAPPINGS
+
+    NODE_DISPLAY_NAME_MAPPINGS = {k: v.TITLE for k, v in NODE_CLASS_MAPPINGS.items()}
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes.py

@@ -0,0 +1,185 @@
+import torch
+import numpy as np
+
+from comfy.ldm.flux.model import (
+    Flux,
+)
+from comfy.ldm.flux.layers import (
+    timestep_embedding,
+)
+from torch import Tensor
+
+def teacache_flux_forward(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor = None,
+        control = None,
+        transformer_options={},
+        attn_mask: Tensor = None,
+    ) -> Tensor:
+        patches_replace = transformer_options.get("patches_replace", {})
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        vec = self.time_in(timestep_embedding(timesteps, 256).to(img.dtype))
+        if self.params.guidance_embed:
+            if guidance is None:
+                raise ValueError("Didn't get guidance strength for guidance distilled model.")
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
+
+        vec = vec + self.vector_in(y[:,:self.params.vec_in_dim])
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        blocks_replace = patches_replace.get("dit", {})
+
+        # enable teacache
+        inp = img.clone()
+        vec_ = vec.clone()
+        img_mod1, _ = self.double_blocks[0].img_mod(vec_)
+        modulated_inp = self.double_blocks[0].img_norm1(inp)
+        modulated_inp = (1 + img_mod1.scale) * modulated_inp + img_mod1.shift
+
+        if self.cnt == 0 or self.cnt == self.steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else: 
+            coefficients = [4.98651651e+02, -2.83781631e+02,  5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
+            rescale_func = np.poly1d(coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+                
+        self.previous_modulated_input = modulated_inp 
+        self.cnt += 1
+
+        if self.cnt == self.steps:
+            self.cnt = 0
+
+        if not should_calc:
+            img += self.previous_residual
+        else:
+            ori_img = img.clone()
+            for i, block in enumerate(self.double_blocks):
+                if ("double_block", i) in blocks_replace:
+                    def block_wrap(args):
+                        out = {}
+                        out["img"], out["txt"] = block(img=args["img"],
+                                                    txt=args["txt"],
+                                                    vec=args["vec"],
+                                                    pe=args["pe"],
+                                                    attn_mask=args.get("attn_mask"))
+                        return out
+
+                    out = blocks_replace[("double_block", i)]({"img": img,
+                                                            "txt": txt,
+                                                            "vec": vec,
+                                                            "pe": pe,
+                                                            "attn_mask": attn_mask},
+                                                            {"original_block": block_wrap})
+                    txt = out["txt"]
+                    img = out["img"]
+                else:
+                    img, txt = block(img=img,
+                                    txt=txt,
+                                    vec=vec,
+                                    pe=pe,
+                                    attn_mask=attn_mask)
+
+                if control is not None: # Controlnet
+                    control_i = control.get("input")
+                    if i < len(control_i):
+                        add = control_i[i]
+                        if add is not None:
+                            img += add
+
+            img = torch.cat((txt, img), 1)
+
+            for i, block in enumerate(self.single_blocks):
+                if ("single_block", i) in blocks_replace:
+                    def block_wrap(args):
+                        out = {}
+                        out["img"] = block(args["img"],
+                                        vec=args["vec"],
+                                        pe=args["pe"],
+                                        attn_mask=args.get("attn_mask"))
+                        return out
+
+                    out = blocks_replace[("single_block", i)]({"img": img,
+                                                            "vec": vec,
+                                                            "pe": pe,
+                                                            "attn_mask": attn_mask}, 
+                                                            {"original_block": block_wrap})
+                    img = out["img"]
+                else:
+                    img = block(img, vec=vec, pe=pe, attn_mask=attn_mask)
+
+                if control is not None: # Controlnet
+                    control_o = control.get("output")
+                    if i < len(control_o):
+                        add = control_o[i]
+                        if add is not None:
+                            img[:, txt.shape[1] :, ...] += add
+
+            img = img[:, txt.shape[1] :, ...]
+            self.previous_residual = img - ori_img
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        return img
+
+class TeaCacheForImgGen:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", {"tooltip": "The diffusion model the TeaCache will be applied to."}),
+                "enable_teacache": ("BOOLEAN", {"default": True, "tooltip": "Enable teacache will speed up inference but may lose visual quality."}),
+                "model_type": (["flux"],),
+                "rel_l1_thresh": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "How strongly to cache the output of diffusion model. This value must be non-negative."}),
+                "steps": ("INT", {"default": 25, "min": 1, "max": 10000, "step": 1}),
+            }
+        }
+    
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply_teacache"
+    CATEGORY = "TeaCache"
+    TITLE = "TeaCache For Img Gen"
+    
+    def apply_teacache(self, model, enable_teacache: bool, model_type: str, rel_l1_thresh: float, steps: int):
+        if enable_teacache:
+            if model_type == "flux":
+                model.model.diffusion_model.__class__.cnt = 0
+                model.model.diffusion_model.__class__.rel_l1_thresh = rel_l1_thresh
+                model.model.diffusion_model.__class__.steps = steps
+                model.model.diffusion_model.forward_orig = teacache_flux_forward.__get__(
+                                                        model.model.diffusion_model,
+                                                        model.model.diffusion_model.__class__
+                                                        )
+            else:
+                raise ValueError(f"Unknown type {model_type}")
+        else:
+            if model_type == "flux":
+                model.model.diffusion_model.forward_orig = Flux.forward_orig.__get__(
+                                                        model.model.diffusion_model,
+                                                        model.model.diffusion_model.__class__
+                                                        )
+            else:
+                raise ValueError(f"Unknown type {model_type}")
+
+        return (model,)
+
+NODE_CLASS_MAPPINGS = {
+    "TeaCacheForImgGen": TeaCacheForImgGen,
+}