32 lines
1.6 KiB
Python
32 lines
1.6 KiB
Python
|
"""Capture block-0 QKV before and after ComfyUI attention preparation."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||
|
|
source = model.read_text(encoding="utf-8")
|
||
|
|
|
||
|
|
old = " block._h3_capture_index = i\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
||
|
|
new = " block._h3_capture_index = i\n block.attn._h3_capture_index = i\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
||
|
|
if source.count(old) == 1:
|
||
|
|
source = source.replace(old, new)
|
||
|
|
elif new not in source:
|
||
|
|
raise RuntimeError("Unable to locate the H3 block loop.")
|
||
|
|
|
||
|
|
old = " v = v.transpose(0, 1).unsqueeze(0)\n out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)\n"
|
||
|
|
new = (
|
||
|
|
" v = v.transpose(0, 1).unsqueeze(0)\n"
|
||
|
|
" if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE:\n"
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" torch.save({\"q\": q.detach().cpu(), \"k\": k.detach().cpu(), \"v\": v.detach().cpu()}, os.path.join(capture_dir, \"block0_qkv_prepared.pt\"))\n"
|
||
|
|
" out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)\n"
|
||
|
|
)
|
||
|
|
if source.count(old) == 1:
|
||
|
|
source = source.replace(old, new)
|
||
|
|
elif new not in source:
|
||
|
|
raise RuntimeError("Unable to locate the prepared QKV call.")
|
||
|
|
|
||
|
|
model.write_text(source, encoding="utf-8")
|
||
|
|
print("Applied H3 block-0 QKV capture patch.")
|