29 lines
1.3 KiB
Python
29 lines
1.3 KiB
Python
"""Add only per-block output capture to the minimal H3 reference hook."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
|
|
old = " for i, block in enumerate(self.blocks):\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
|
new = (
|
|
" for i, block in enumerate(self.blocks):\n"
|
|
" comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate H3 block loop.")
|
|
|
|
old = " h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
|
|
new = (
|
|
" 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 source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate H3 direct block execution.")
|
|
source = source.replace(old, new)
|
|
model.write_text(source, encoding="utf-8")
|
|
print("Applied minimal H3 per-block capture patch.")
|