32 lines
1.8 KiB
Python
32 lines
1.8 KiB
Python
"""Capture MiniMax Qwen layer-50 output before H3 token refinement."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
path = Path("/opt/ComfyUI/comfy/text_encoders/minimax.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")
|
|
old = (
|
|
" return super().forward(input_ids, attention_mask=attention_mask, embeds=embeds,\n"
|
|
" num_tokens=num_tokens, intermediate_output=intermediate_output,\n"
|
|
" final_layer_norm_intermediate=final_layer_norm_intermediate,\n"
|
|
" dtype=dtype, embeds_info=embeds_info, **kwargs)\n"
|
|
)
|
|
new = (
|
|
" output = super().forward(input_ids, attention_mask=attention_mask, embeds=embeds,\n"
|
|
" num_tokens=num_tokens, intermediate_output=intermediate_output,\n"
|
|
" final_layer_norm_intermediate=final_layer_norm_intermediate,\n"
|
|
" dtype=dtype, embeds_info=embeds_info, **kwargs)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
|
" if capture_dir:\n"
|
|
" os.makedirs(capture_dir, exist_ok=True)\n"
|
|
" torch.save(input_ids.detach().cpu() if input_ids is not None else torch.empty(0, dtype=torch.long), os.path.join(capture_dir, \"qwen_input_ids.pt\"))\n"
|
|
" torch.save(embeds.detach().cpu(), os.path.join(capture_dir, \"qwen_input_embeds.pt\"))\n"
|
|
" torch.save(output[0].detach().cpu(), os.path.join(capture_dir, \"qwen_layer50.pt\"))\n"
|
|
" return output\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate MiniMax Qwen forward return.")
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied MiniMax Qwen layer-50 capture patch.")
|