33 lines
1.7 KiB
Python
33 lines
1.7 KiB
Python
"""Add only block-0 MLP projection saves to the minimal H3 capture."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
old = (
|
|
" def forward(self, x):\n"
|
|
" return comfy.ops.linear_input_act(self.fc2, self.fc1(x), \"swiglu\")\n"
|
|
)
|
|
new = (
|
|
" def forward(self, x):\n"
|
|
" fc1 = self.fc1(x)\n"
|
|
" activated = comfy.ops.INPUT_ACT_EAGER[\"swiglu\"](fc1)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
|
|
" if capture_dir:\n"
|
|
" torch.save(fc1.detach().cpu(), os.path.join(capture_dir, \"block0_mlp_fc1.pt\"))\n"
|
|
" torch.save(activated.detach().cpu(), os.path.join(capture_dir, \"block0_mlp_activated.pt\"))\n"
|
|
" output = self.fc2(activated)\n"
|
|
" if capture_dir: torch.save(output.detach().cpu(), os.path.join(capture_dir, \"block0_mlp_fc2.pt\"))\n"
|
|
" return output\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate pristine H3 MLP.forward.")
|
|
source = source.replace(old, new)
|
|
|
|
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.mlp._h3_capture_index = i\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate H3 block index assignment.")
|
|
model.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied minimal H3 block-0 MLP projection capture patch.")
|