"""Capture one complete 50-layer MiniMax Qwen text trace for offline parity work."""

from pathlib import Path


path = Path("/opt/ComfyUI/comfy/text_encoders/llama.py")
source = path.read_text(encoding="utf-8")
if "import os\n" not in source:
    source = "import os\n" + source
old = (
    "            x, current_kv = layer(\n"
    "                x=x,\n"
    "                attention_mask=mask,\n"
    "                freqs_cis=freqs_cis,\n"
    "                optimized_attention=optimized_attention,\n"
    "                past_key_value=past_kv,\n"
    "            )\n"
)
new = old + (
    "            capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
    "            if capture_dir and x.shape[-1] == 5120 and len(self.layers) == 50:\n"
    "                layer_dir = os.path.join(capture_dir, \"qwen_layers\")\n"
    "                os.makedirs(layer_dir, exist_ok=True)\n"
    "                torch.save(x.detach().cpu(), os.path.join(layer_dir, f\"{i:02d}.pt\"))\n"
)
if source.count(old) != 1:
    raise RuntimeError("Unable to locate Qwen decoder layer loop.")
path.write_text(source.replace(old, new), encoding="utf-8")
print("Applied complete MiniMax Qwen layer trace patch.")