"""Capture block-0 RMSNorm and AdaLN components before modulation."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")
old = (
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
)
new = (
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
    "        norm1 = self.norm1(x)\n"
    "        if capture_dir: torch.save({\"norm\": norm1.detach().cpu(), \"shift\": shift_msa.detach().cpu(), \"scale\": scale_msa.detach().cpu()}, os.path.join(capture_dir, \"block0_norm1_adaln.pt\"))\n"
    "        h = _mod_scale_shift(norm1, shift_msa, scale_msa, mod_segments)\n"
)
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate block-0 norm1 capture point.")
model.write_text(source, encoding="utf-8")
print("Applied H3 block-0 RMSNorm/AdaLN capture patch.")