42 lines
2.7 KiB
Python
42 lines
2.7 KiB
Python
"""Add only block-0 sublayer captures to the minimal H3 reference hook."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
|
|
old = " for i, block in enumerate(self.blocks):\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
|
new = " for i, block in enumerate(self.blocks):\n block._h3_capture_index = i\n comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)\n"
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate pristine H3 block loop.")
|
|
source = source.replace(old, new)
|
|
|
|
old = (
|
|
" shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
|
|
" h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
|
|
" x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)\n"
|
|
" h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
|
|
" return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)\n"
|
|
)
|
|
new = (
|
|
" shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
|
|
" h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)\n"
|
|
" if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm1.pt\"))\n"
|
|
" attention = self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options)\n"
|
|
" if capture_dir: torch.save(attention.detach().cpu(), os.path.join(capture_dir, \"block0_attention.pt\"))\n"
|
|
" x = _mod_gate(x, gate_msa, attention, mod_segments)\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_attention.pt\"))\n"
|
|
" h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)\n"
|
|
" if capture_dir: torch.save(h.detach().cpu(), os.path.join(capture_dir, \"block0_norm2.pt\"))\n"
|
|
" mlp = self.mlp(h)\n"
|
|
" if capture_dir: torch.save(mlp.detach().cpu(), os.path.join(capture_dir, \"block0_mlp.pt\"))\n"
|
|
" x = _mod_gate(x, gate_mlp, mlp, mod_segments)\n"
|
|
" if capture_dir: torch.save(x.detach().cpu(), os.path.join(capture_dir, \"block0_post_mlp.pt\"))\n"
|
|
" return x\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate pristine H3 DiTBlock.forward.")
|
|
model.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied minimal H3 block-0 sublayer capture patch.")
|