28 lines
1.3 KiB
Python
28 lines
1.3 KiB
Python
"""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.")
|