29 lines
1.5 KiB
Python
29 lines
1.5 KiB
Python
"""Capture H3 curve embeddings after every transformer block."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
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"
|
|
" if capture_dir and H3_CAPTURE_ACTIVE and i == 0:\n"
|
|
" torch.save({\"t_emb\": t_emb.clone().detach().cpu(), \"timesteps\": t_vals.clone().detach().cpu()}, os.path.join(capture_dir, \"t_emb_before_block0.pt\"))\n"
|
|
" h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)\n"
|
|
" if capture_dir and H3_CAPTURE_ACTIVE:\n"
|
|
" temb_dir = os.path.join(capture_dir, \"t_emb\")\n"
|
|
" os.makedirs(temb_dir, exist_ok=True)\n"
|
|
" torch.save({\"value\": t_emb.detach().cpu(), \"class\": f\"{type(block.adaln_proj.linear).__module__}.{type(block.adaln_proj.linear).__qualname__}\", \"data_ptr\": t_emb.data_ptr(), \"stride\": t_emb.stride()}, os.path.join(temb_dir, f\"{i:02d}.pt\"))\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 transformer block loop.")
|
|
model.write_text(source, encoding="utf-8")
|
|
print("Applied H3 t_emb trace patch.")
|