38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
"""Extend the H3 capture hook to persist every block output once."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
|
|
old = 'if capture_dir and not H3_CAPTURE_ACTIVE:'
|
|
new = 'if capture_dir and not H3_CAPTURE_ACTIVE and not os.path.exists(os.path.join(capture_dir, "blocks_complete")):'
|
|
if source.count(old) == 1:
|
|
source = source.replace(old, new)
|
|
elif new not in source:
|
|
raise RuntimeError("Unable to locate the H3 input capture condition.")
|
|
|
|
old = (
|
|
' else:\n'
|
|
' h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n'
|
|
' if prefetch_queue is not None:\n'
|
|
)
|
|
new = (
|
|
' else:\n'
|
|
' h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n'
|
|
' if capture_dir and H3_CAPTURE_ACTIVE:\n'
|
|
' block_dir = os.path.join(capture_dir, "blocks")\n'
|
|
' os.makedirs(block_dir, exist_ok=True)\n'
|
|
' torch.save(h.detach().cpu(), os.path.join(block_dir, f"{i:02d}.pt"))\n'
|
|
' if capture_dir and H3_CAPTURE_ACTIVE:\n'
|
|
' open(os.path.join(capture_dir, "blocks_complete"), "a").close()\n'
|
|
' if prefetch_queue is not None:\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.")
|
|
|
|
model.write_text(source, encoding="utf-8")
|
|
print("Applied H3 per-block capture patch.")
|