15 lines
596 B
Python
15 lines
596 B
Python
|
"""Keep block-0 QKV diagnostics from being overwritten by block-2 capture."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||
|
|
source = model.read_text(encoding="utf-8")
|
||
|
|
old = 'getattr(self, "_h3_capture_index", -1) in (0, 2) and H3_CAPTURE_ACTIVE'
|
||
|
|
new = 'getattr(self, "_h3_capture_index", -1) == 0 and H3_CAPTURE_ACTIVE'
|
||
|
|
if source.count(old) < 2:
|
||
|
|
raise RuntimeError("Unable to locate both H3 Attention QKV capture conditions.")
|
||
|
|
source = source.replace(old, new, 2)
|
||
|
|
model.write_text(source, encoding="utf-8")
|
||
|
|
print("Restricted H3 QKV capture to block 0.")
|