"""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 = (
    "        capture_input = t_emb.clone() if H3_CAPTURE_ACTIVE else None\n"
    "        adaln_output = self.adaln_proj(t_emb)\n"
    "        capture_replay = self.adaln_proj.linear(capture_input.clone()) if H3_CAPTURE_ACTIVE else None\n"
    "        capture_forward = self.adaln_proj.linear._forward(capture_input.clone(), self.adaln_proj.linear.weight, self.adaln_proj.linear.bias) if H3_CAPTURE_ACTIVE else None\n"
    "        shift, scale = adaln_output\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.cuda.synchronize(x.device)\n"
    "            linear = self.adaln_proj.linear\n"
    "            torch.save({\"hidden\": x.detach().cpu(), \"adaln_input\": capture_input.detach().cpu(), \"adaln_output\": torch.cat((shift, scale), dim=-1).detach().cpu(), \"adaln_replay\": capture_replay.detach().cpu(), \"adaln_forward\": capture_forward.detach().cpu(), \"t_emb\": t_emb.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"adaln_weight\": linear.weight.detach().cpu(), \"adaln_bias\": linear.bias.detach().cpu(), \"adaln_metadata\": {\"class\": f\"{type(linear).__module__}.{type(linear).__qualname__}\", \"weight_class\": f\"{type(linear.weight).__module__}.{type(linear.weight).__qualname__}\", \"quant_format\": getattr(linear, \"quant_format\", None), \"layout_type\": getattr(linear, \"layout_type\", None), \"full_precision_mm\": getattr(linear, \"_full_precision_mm\", None), \"input_scale\": getattr(linear, \"input_scale\", None)}, \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video_weight\": self.video_out.weight.detach().cpu(), \"video_bias\": self.video_out.bias.detach().cpu(), \"audio_weight\": self.audio_out.weight.detach().cpu(), \"audio_bias\": self.audio_out.bias.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, 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.")