"""Capture H3 text states before and after the token refiner."""

from pathlib import Path


model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
source = model.read_text(encoding="utf-8")
old = (
    "        text_states = context[0]\n"
    "        if text_states.shape[-1] != self.hidden_size:\n"
    "            text_states = self.token_refiner(self.condition_proj(text_states),\n"
    "                                             transformer_options=transformer_options)\n"
)
new = (
    "        text_states = context[0]\n"
    "        capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
    "        if capture_dir:\n"
    "            torch.save(text_states.detach().cpu(), os.path.join(capture_dir, \"text_qwen.pt\"))\n"
    "        if text_states.shape[-1] != self.hidden_size:\n"
    "            text_states = self.token_refiner(self.condition_proj(text_states),\n"
    "                                             transformer_options=transformer_options)\n"
    "        if capture_dir:\n"
    "            torch.save(text_states.detach().cpu(), os.path.join(capture_dir, \"text_refined.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate H3 text-refiner path.")
model.write_text(source.replace(old, new), encoding="utf-8")
print("Applied H3 text-state capture patch.")