29 lines
1.2 KiB
Python
29 lines
1.2 KiB
Python
"""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.")
|