"""Capture Comfy H3 FP32 patch-projection inputs, outputs, and matmul policy."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = path.read_text(encoding="utf-8")
old = (
    "        video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
    "        audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
)
new = (
    "        video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
    "        audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
    "        if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_patch_projections.pt\")):\n"
    "            torch.save({\"video_rows\": all_video_rows.detach().cpu(), \"audio_rows\": all_audio_rows.detach().cpu(), \"video_embed_fp32\": video_embed_fp32.detach().cpu(), \"audio_embed_fp32\": audio_embed_fp32.detach().cpu(), \"video_weight\": self.video_patch_proj.weight.detach().cpu(), \"video_bias\": self.video_patch_proj.bias.detach().cpu(), \"audio_weight\": self.audio_patch_proj.weight.detach().cpu(), \"audio_bias\": self.audio_patch_proj.bias.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"h3_patch_projections.pt\"))\n"
    "        video_embed = video_embed_fp32.to(dtype)\n"
    "        audio_embed = audio_embed_fp32.to(dtype)\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate H3 patch-projection calls.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied H3 patch-projection probe.")