17 lines
657 B
Python
17 lines
657 B
Python
"""Extend the local H3 capture hook to persist block-2 intermediates."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
|
|
source = source.replace(
|
|
'getattr(self, "_h3_capture_index", -1) == 0 and H3_CAPTURE_ACTIVE',
|
|
'getattr(self, "_h3_capture_index", -1) in (0, 2) and H3_CAPTURE_ACTIVE',
|
|
)
|
|
for name in ("norm1", "attention", "post_attention", "norm2", "mlp", "post_mlp"):
|
|
source = source.replace(f'"block0_{name}.pt"', f'f"block{{self._h3_capture_index}}_{name}.pt"')
|
|
|
|
model.write_text(source, encoding="utf-8")
|
|
print("Applied H3 block-2 sublayer capture patch.")
|