"""Capture Comfy's loaded block-0 NVFP4 fc1 tensor layout once."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")
old = "        fc1 = self.fc1(x)\n"
new = (
    "        fc1 = self.fc1(x)\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"
    "            weight = self.fc1.weight\n"
    "            torch.save({\"qdata\": weight._qdata.detach().cpu(), \"scale\": weight._params.scale.detach().cpu(), \"block_scale\": weight._params.block_scale.detach().cpu(), \"orig_dtype\": str(weight._params.orig_dtype), \"orig_shape\": weight._params.orig_shape}, os.path.join(capture_dir, \"block0_mlp_fc1_weight.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate minimal H3 MLP fc1 assignment.")
model.write_text(source.replace(old, new), encoding="utf-8")
print("Applied H3 block-0 fc1 weight capture patch.")