"""Capture H3 final-layer intermediates from one reference inference."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")
old = (
    "        shift, scale = self.adaln_proj(t_emb)\n"
    "        va, vb, vrow = video_seg\n"
    "        aa, ab, arow = audio_seg\n"
    "        hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
    "        ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
    "        return self.video_out(hv), self.audio_out(ha)\n"
)
new = (
    "        shift, scale = self.adaln_proj(t_emb)\n"
    "        va, vb, vrow = video_seg\n"
    "        aa, ab, arow = audio_seg\n"
    "        norm_v = self.norm(x[va:vb])\n"
    "        norm_a = self.norm(x[aa:ab])\n"
    "        hv = (norm_v * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
    "        ha = (norm_a * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
    "        video = self.video_out(hv)\n"
    "        audio = self.audio_out(ha)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if H3_CAPTURE_ACTIVE else None\n"
    "        if capture_dir:\n"
    "            torch.save({\"hidden\": x.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu()}, os.path.join(capture_dir, \"final.pt\"))\n"
    "        return video, audio\n"
)
if source.count(old) == 1:
    source = source.replace(old, new)
elif new not in source:
    raise RuntimeError("Unable to locate H3 FinalLayer.forward.")
model.write_text(source, encoding="utf-8")
print("Applied H3 final-layer capture patch.")