From 35231722cb91bd0e4b3a10f1d42d6ae8b8e506ab Mon Sep 17 00:00:00 2001 From: YunjieYu Date: Tue, 7 Jan 2025 17:06:20 +0800 Subject: [PATCH] first commit --- README.md | 10 ++ __init__.py | 10 ++ __pycache__/__init__.cpython-310.pyc | Bin 0 -> 500 bytes __pycache__/nodes.cpython-310.pyc | Bin 0 -> 4567 bytes nodes.py | 185 +++++++++++++++++++++++++++ 5 files changed, 205 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 __pycache__/__init__.cpython-310.pyc create mode 100644 __pycache__/nodes.cpython-310.pyc create mode 100644 nodes.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..e12ba50 --- /dev/null +++ b/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!!! \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..aefb566 --- /dev/null +++ b/__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"] diff --git a/__pycache__/__init__.cpython-310.pyc b/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98dd9e02d591c6e6c891b7afa9b4ab2a77c84bc2 GIT binary patch literal 500 zcmYjN%}&BV5T4z(PyrL87vdXmkoo|GNT4xI07;555Y2{k7p%4{_D55Nlb*bI@CEef zEA7=2ui!;@A*eH%FEg_{-+r@vd%FPAo?jj&76SNG#eXwpaU>&dOSsQbpeTSLI)iCJ z2Zpm91j+}p@4_8WOtnA@^`VY{>eP4wroS2Q(1J%a%+VapPw+6W?pDmDWQ6Wgz6up2 z_PQr6n>5=E*Cm}sujh1+T~D2DOp}Q&qbVaUOYmV17v?!$paq_zIlh8de1l#1P;^G$ z>9_5QCUEsoplg9{BQ?B&FWb80A+xeHOQ%;pvpi0j)#TwgJ9R8CjT0V{sqDp@=0Ts0 znj`O$S%mmgKOv+#%S53@eJ}BNIIC|dQCLD5BK4t>QY-dXJGC6Q*KQ0*x6!eGSF4z! zC@05RHBI~=7RIU&MJJr`D6yl6N5WWXSW_8(!ooNzD%w0DqaYyUZ=-dlLqF;Vc6By_u8ydrX(C4C;YO7hk15315^ zfB}eq4V)GR5YpGS$VMsEmvu^3jCeFPvJsz12sJkKpXjMJFj9R>-;g6Su~Lhgn`Dr~ zo&1(uFQ?kFy&Gs8b5-Ymb!_ZK@m6dyKhWsxFh#UX%0rv-h^2cjq;nMO=h$$ zU7gSdmDB)zBdIQrq-Kgenxpw^WTS$+qiF?u)wBdDY$e4=st?Atw2jfUtbWHpbDS1t zw85TqY)gNEY>aR0*;$)NM{$3G=F;)y$tf+Jn9wlVi!poNpBwmg$grINzzv^F@fr40^lN!ezH67F5m_}^%GV=}aEzalBQvx@E%&8hQ}AH@^+`AG3Hl}^SVr?AfQ zeUEAzdMLK_Fr9o0YvY^xY3;=Zq@Ve@Ml0z|o1`PF+Anl|?@zV1uAJB2*q4rEdovm+ z8L&<@Ez;5J79C6X`T8c=*q`pF<1^ZZf!{s&)x*-ycAdRu;SE?Twc!Oi=s%g-u+0QC zGn>vNhtkCe^_R-q)qA;hDOU;eI8<^V>2^-XX)Q zN`}1%z)S|Ur)0<^n7VH6k|5*UR-9)irmvas>>n+cY5U^{6zal%>HNtT>;#BZ-sKTiB0VASbGlm+#hxUV5S86YZA!sSAv zW_4bD=G|KNq1R4ye*3RC-g)CMmp|xy=>GYuS3myYiO#W?K6>G=?^HkNeDtT^z5U?L zzaKv%$~5XP2FzUyqSi`F(inaO1DTy<@{W%J`OprmgbW33B55mZC~Qg$&$703nv-4! z4eLkYibfD;;^E@M4d$KlV8Fv_ zx!zk7MMccL!oA*8csCveGId+esXtpEzhzqMwS;*sdCzU5R=BXt6_3;UQ7Fv|;$!vO zH24(mi)zIEjvsnKrdK7kl2TZ{bQMvKQC+@K?NUif?R5;kS?y4`358qN%%QrW=-jrV z@6x1kmumSmNb`LX9gs-YBJN0ZFHjErdcyHQ|+eB@5e!rgKZ#Lz4joV zk#jSYo>_dx3;Jx1b5zirDxM+ACgQ;LLoQ=SkR%coP1br$#j zLs5Y5BoW7bQN$3xu6!t`YH=tm?6sDJ(ee^e>ap!o&bY8zK@>7kxzpd{te_*JL_sna zM!b7OSTT&s@U!H%6Hp2hjL{av&HY}h)$jHLFJTlJAaH{tcsx#2u_s{<5=6n;5!#&Y zaklD5{n!QV?em%84R0|PW|w&(DEIop^b^(182*j7+^CnxWTirpOK(s#hONvj@gIQ7l?eoT(U=C7f3y01q^OSr%EUHoeL^1v>?1M=Y^;7? zFC!_K1gU{kaYl~I*csO`)5`^-XMg6coMnD)=Ze@Lt`%<;F|UGB5kOuaCF7EYUbeUI ze33KFTSaS>6!B7a?vQT2KC8Gfb^aqT+ydUrQ)1A>s&!`bQYQ8 z^?HGiDw#xek$?RB`Ezq87aHrQ=Ryxv4plizNO(rl(oHnknrarx-exi5Z0M<=T9~|@4YL(e;n^%M}6CaLnk_~ zO#JceKi97wU%&K3bVb(p2)Ta^EIO*DC_%F9=|t_j#jHDxB|mmnWggAC2}l`6;lm;8 zc!|Ht>NShMh!=&?MNTxj-ph!4q^pcaV}7A=Qzm3S2IP-EfUdz*=lPH2swdGRP|Xq+ z!juJKgehgYLuHAgT6blR=BMTPA4nuC1{x4O&Z2IofQZuk!i7r>xAEMCxr^`MkSeLf zKvW@m#t~EI=jK=I2ZpBGakTm!_i8$C;;N`hd#>%U(zJ)S)+dG|2Lv3rssmr;ZONI` z{B1^CGW88VIW_^xf8WC(CHy-n0tT&njiPN?)|T~U;u`qpW;ixeYiDdFxj{ag4cE4| zZ@D&FWO#B?phjx^Qu`HA|8Qlr&m_fV`~SOHGZ*+v;7~Po>y=C4N*G-U9ccxpR^kj> zVq_;KWUhtn3a;Ay9Lsx+CO5JJ8r2Kj=R j 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, +}