19 lines
1 KiB
Python
19 lines
1 KiB
Python
"""Capture Comfy's loaded block-0 NVFP4 fc1 tensor layout once."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
source = model.read_text(encoding="utf-8")
|
|
old = " fc1 = self.fc1(x)\n"
|
|
new = (
|
|
" fc1 = self.fc1(x)\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\") if getattr(self, \"_h3_capture_index\", -1) == 0 and H3_CAPTURE_ACTIVE else None\n"
|
|
" if capture_dir:\n"
|
|
" weight = self.fc1.weight\n"
|
|
" torch.save({\"qdata\": weight._qdata.detach().cpu(), \"scale\": weight._params.scale.detach().cpu(), \"block_scale\": weight._params.block_scale.detach().cpu(), \"orig_dtype\": str(weight._params.orig_dtype), \"orig_shape\": weight._params.orig_shape}, os.path.join(capture_dir, \"block0_mlp_fc1_weight.pt\"))\n"
|
|
)
|
|
if source.count(old) != 1:
|
|
raise RuntimeError("Unable to locate minimal H3 MLP fc1 assignment.")
|
|
model.write_text(source.replace(old, new), encoding="utf-8")
|
|
print("Applied H3 block-0 fc1 weight capture patch.")
|