"""Probe loaded Comfy H3 patch projections on their captured FP32 rows."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = path.read_text(encoding="utf-8")
if "import os\n" not in source:
    source = source.replace("import math\n", "import math\nimport os\n", 1)
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_loaded_patch_projection_probe.pt\")):\n"
    "            def probe(layer, rows, output):\n"
    "                weight, bias, state = comfy.ops.cast_bias_weight(layer, rows, offloadable=True)\n"
    "                try:\n"
    "                    cast_output = torch.nn.functional.linear(rows, weight, bias)\n"
    "                    return {\"rows\": rows.detach().cpu(), \"output\": output.detach().cpu(), \"forward_cast\": layer.forward_comfy_cast_weights(rows).detach().cpu(), \"effective_weight\": weight.detach().cpu(), \"effective_bias\": None if bias is None else bias.detach().cpu(), \"effective_weight_dtype\": str(weight.dtype), \"effective_bias_dtype\": None if bias is None else str(bias.dtype), \"cast_output\": cast_output.detach().cpu(), \"rows_dtype\": str(rows.dtype), \"rows_stride\": tuple(rows.stride()), \"weight_dtype\": str(layer.weight.dtype), \"weight_stride\": tuple(layer.weight.stride()), \"bias_dtype\": None if layer.bias is None else str(layer.bias.dtype), \"bias_stride\": None if layer.bias is None else tuple(layer.bias.stride()), \"autocast\": torch.is_autocast_enabled(), \"force_cast\": getattr(layer, \"comfy_force_cast_weights\", False), \"weight_functions\": len(layer.weight_function), \"bias_functions\": len(layer.bias_function)}\n"
    "                finally:\n"
    "                    comfy.ops.uncast_bias_weight(layer, weight, bias, state)\n"
    "            torch.save({\"video\": probe(self.video_patch_proj, all_video_rows, video_embed_fp32), \"audio\": probe(self.audio_patch_proj, all_audio_rows, audio_embed_fp32)}, os.path.join(capture_dir, \"h3_loaded_patch_projection_probe.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 loaded H3 patch-projection probe.")