34 lines
1.5 KiB
Python
34 lines
1.5 KiB
Python
"""Capture Comfy H3 token-refiner boundaries once for direct parity replay."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.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", 1)
|
|
|
|
old = (
|
|
" def forward(self, x, transformer_options={}):\n"
|
|
" for block in self.blocks:\n"
|
|
" x = block(x, transformer_options=transformer_options)\n"
|
|
" return self.final_norm(x)\n"
|
|
)
|
|
new = (
|
|
" def forward(self, x, transformer_options={}):\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
|
" if capture_dir:\n"
|
|
" torch.save(x.detach().cpu(), os.path.join(capture_dir, \"refiner_input.pt\"))\n"
|
|
" for index, block in enumerate(self.blocks):\n"
|
|
" x = block(x, transformer_options=transformer_options)\n"
|
|
" if capture_dir:\n"
|
|
" torch.save(x.detach().cpu(), os.path.join(capture_dir, f\"refiner_block{index}.pt\"))\n"
|
|
" output = self.final_norm(x)\n"
|
|
" if capture_dir:\n"
|
|
" torch.save(output.detach().cpu(), os.path.join(capture_dir, \"refiner_output.pt\"))\n"
|
|
" return output\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate TokenRefiner.forward.")
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied H3 token-refiner trace patch.")
|