25 lines
1.6 KiB
Python
25 lines
1.6 KiB
Python
|
"""Capture Comfy H3 FP32 patch-projection inputs, outputs, and matmul policy."""
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
path = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
||
|
|
source = path.read_text(encoding="utf-8")
|
||
|
|
old = (
|
||
|
|
" video_embed = self.video_patch_proj(all_video_rows).to(dtype)\n"
|
||
|
|
" audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)\n"
|
||
|
|
)
|
||
|
|
new = (
|
||
|
|
" video_embed_fp32 = self.video_patch_proj(all_video_rows)\n"
|
||
|
|
" audio_embed_fp32 = self.audio_patch_proj(all_audio_rows)\n"
|
||
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
||
|
|
" if capture_dir and not os.path.exists(os.path.join(capture_dir, \"h3_patch_projections.pt\")):\n"
|
||
|
|
" torch.save({\"video_rows\": all_video_rows.detach().cpu(), \"audio_rows\": all_audio_rows.detach().cpu(), \"video_embed_fp32\": video_embed_fp32.detach().cpu(), \"audio_embed_fp32\": audio_embed_fp32.detach().cpu(), \"video_weight\": self.video_patch_proj.weight.detach().cpu(), \"video_bias\": self.video_patch_proj.bias.detach().cpu(), \"audio_weight\": self.audio_patch_proj.weight.detach().cpu(), \"audio_bias\": self.audio_patch_proj.bias.detach().cpu(), \"matmul\": {\"allow_tf32\": torch.backends.cuda.matmul.allow_tf32, \"precision\": torch.get_float32_matmul_precision()}}, os.path.join(capture_dir, \"h3_patch_projections.pt\"))\n"
|
||
|
|
" video_embed = video_embed_fp32.to(dtype)\n"
|
||
|
|
" audio_embed = audio_embed_fp32.to(dtype)\n"
|
||
|
|
)
|
||
|
|
if source.count(old) != 1:
|
||
|
|
raise RuntimeError("Unable to locate H3 patch-projection calls.")
|
||
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
||
|
|
print("Applied H3 patch-projection probe.")
|