37 lines
1.8 KiB
Python
37 lines
1.8 KiB
Python
|
"""Capture H3 final-layer intermediates from one reference inference."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||
|
|
source = model.read_text(encoding="utf-8")
|
||
|
|
old = (
|
||
|
|
" shift, scale = self.adaln_proj(t_emb)\n"
|
||
|
|
" va, vb, vrow = video_seg\n"
|
||
|
|
" aa, ab, arow = audio_seg\n"
|
||
|
|
" hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
|
||
|
|
" ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
|
||
|
|
" return self.video_out(hv), self.audio_out(ha)\n"
|
||
|
|
)
|
||
|
|
new = (
|
||
|
|
" shift, scale = self.adaln_proj(t_emb)\n"
|
||
|
|
" va, vb, vrow = video_seg\n"
|
||
|
|
" aa, ab, arow = audio_seg\n"
|
||
|
|
" norm_v = self.norm(x[va:vb])\n"
|
||
|
|
" norm_a = self.norm(x[aa:ab])\n"
|
||
|
|
" hv = (norm_v * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32)\n"
|
||
|
|
" ha = (norm_a * (1.0 + scale[arow]) + shift[arow]).to(torch.float32)\n"
|
||
|
|
" video = self.video_out(hv)\n"
|
||
|
|
" audio = self.audio_out(ha)\n"
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if H3_CAPTURE_ACTIVE else None\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" torch.save({\"hidden\": x.detach().cpu(), \"norm_v\": norm_v.detach().cpu(), \"norm_a\": norm_a.detach().cpu(), \"shift\": shift.detach().cpu(), \"scale\": scale.detach().cpu(), \"video_hidden\": hv.detach().cpu(), \"audio_hidden\": ha.detach().cpu(), \"video\": video.detach().cpu(), \"audio\": audio.detach().cpu()}, os.path.join(capture_dir, \"final.pt\"))\n"
|
||
|
|
" return video, audio\n"
|
||
|
|
)
|
||
|
|
if source.count(old) == 1:
|
||
|
|
source = source.replace(old, new)
|
||
|
|
elif new not in source:
|
||
|
|
raise RuntimeError("Unable to locate H3 FinalLayer.forward.")
|
||
|
|
model.write_text(source, encoding="utf-8")
|
||
|
|
print("Applied H3 final-layer capture patch.")
|