"""Capture block-0 AdaLN, attention, and MLP intermediates once."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")

old = "        for i, block in enumerate(self.blocks):\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
new = "        for i, block in enumerate(self.blocks):\n            block._h3_capture_index = i\n            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate the H3 block loop.")

old = (
    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
    "        h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
    "        x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)\n"
    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
    "        return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)\n"
)
new = (
    "        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
    "        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"
    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm1.pt\"))\n"
    "        attn_out = self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options)\n"
    "        if capture_dir: torch.save(attn_out.detach().cpu(), os.path.join(capture_dir, \"block0_attention.pt\"))\n"
    "        x = _mod_gate(x, gate_msa, attn_out, mod_segments)\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_attention.pt\"))\n"
    "        h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
    "        if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm2.pt\"))\n"
    "        mlp_out = self.mlp(h)\n"
    "        if capture_dir: torch.save(mlp_out.detach().cpu(), os.path.join(capture_dir, \"block0_mlp.pt\"))\n"
    "        x = _mod_gate(x, gate_mlp, mlp_out, mod_segments)\n"
    "        if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_mlp.pt\"))\n"
    "        return x\n"
)
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate H3 DiTBlock.forward.")

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