wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-26 01:39:25 +08:00
Code Issues Actions 28 Packages Projects Releases Wiki Activity

Don't crash when using flux kv cache with split batches.

Browse Source
This commit is contained in:
comfyanonymous 2026-06-11 15:36:33 -07:00 committed by GitHub
parent fb991e2c1e
commit 4cc7ada854
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy_extras/nodes_flux.py
View File

@ -245,6 +245,11 @@ class KV_Attn_Input:
cache_key = "{}_{}".format(extra_options["block_type"], extra_options["block_index"]) cache_key = "{}_{}".format(extra_options["block_type"], extra_options["block_index"])
if cache_key in self.cache: if cache_key in self.cache:
kk, vv = self.cache[cache_key] kk, vv = self.cache[cache_key]
# Fix batch size changing.
kk = comfy.utils.repeat_to_batch_size(kk, k.shape[0])
vv = comfy.utils.repeat_to_batch_size(vv, v.shape[0])
self.set_cache = False self.set_cache = False
return {"q": q, "k": torch.cat((k, kk), dim=2), "v": torch.cat((v, vv), dim=2)} return {"q": q, "k": torch.cat((k, kk), dim=2), "v": torch.cat((v, vv), dim=2)}