"""Capture block-0 MLP projections for direct NVFP4 parity diagnostics."""

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:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate H3 MLP.forward.")

old = "            block._h3_capture_index = i\n            block.attn._h3_capture_index = i\n"
new = "            block._h3_capture_index = i\n            block.attn._h3_capture_index = i\n            block.mlp._h3_capture_index = i\n"
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate H3 capture block-index assignment.")

model.write_text(source, encoding="utf-8")
print("Applied H3 block-0 MLP projection capture patch.")