39 lines
1.9 KiB
Python
39 lines
1.9 KiB
Python
"""Instrument ComfyUI v0.31.1 H3 once to capture direct-runner inputs/outputs."""
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
def replace_once(path, old, new):
|
|
source = path.read_text(encoding="utf-8")
|
|
if new in source:
|
|
return
|
|
if source.count(old) != 1:
|
|
raise RuntimeError(f"Expected one matching block in {path}.")
|
|
path.write_text(source.replace(old, new), encoding="utf-8")
|
|
|
|
|
|
model = Path("/opt/ComfyUI/comfy/ldm/minimax/model.py")
|
|
replace_once(model, "import math\n", "import math\nimport os\n")
|
|
replace_once(model, "VISUAL_COND_TIMESTEP = 0.999\n", "VISUAL_COND_TIMESTEP = 0.999\nH3_CAPTURE_ACTIVE = False\n")
|
|
replace_once(
|
|
model,
|
|
" # blocks\n patches_replace = transformer_options.get(\"patches_replace\", {})\n",
|
|
" # Capture a single fully assembled payload before the first transformer block.\n"
|
|
" global H3_CAPTURE_ACTIVE\n"
|
|
" capture_dir = os.getenv(\"H3_CAPTURE_DIR\")\n"
|
|
" if capture_dir and not H3_CAPTURE_ACTIVE:\n"
|
|
" H3_CAPTURE_ACTIVE = True\n"
|
|
" os.makedirs(capture_dir, exist_ok=True)\n"
|
|
" torch.save({\"hidden\": h.detach().cpu(), \"timesteps\": t_vals.detach().cpu(), \"position_ids\": layout.position_ids, \"segments\": mod_segments}, os.path.join(capture_dir, \"input.pt\"))\n\n"
|
|
" # blocks\n patches_replace = transformer_options.get(\"patches_replace\", {})\n",
|
|
)
|
|
replace_once(
|
|
model,
|
|
" return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
|
|
" if capture_dir and H3_CAPTURE_ACTIVE:\n"
|
|
" torch.save({\"video\": video_out.detach().cpu(), \"audio\": audio_out.detach().cpu(), \"video_segment\": video_seg, \"audio_segment\": audio_seg}, os.path.join(capture_dir, \"output.pt\"))\n"
|
|
" H3_CAPTURE_ACTIVE = False\n"
|
|
" return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]\n",
|
|
)
|
|
|
|
print("Applied H3 direct-runner capture patch.")
|