daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Capture Comfy denoised wrapper boundary

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 16:23:08 +07:00
parent 58791dbd61
commit 1e6a65489a
1 changed files with 40 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
tools/patch_comfy_wrapper_capture.py Normal file
View file

@ -0,0 +1,40 @@
"""Capture Comfy sampler model-wrapper inputs and denoised outputs."""
from pathlib import Path
path = Path("/opt/ComfyUI/comfy/samplers.py")
source = path.read_text(encoding="utf-8")
if "import os\n" not in source:
source = source.replace("import logging\n", "import logging\nimport os\n")
if "H3_WRAPPER_CAPTURE_INDEX = 0\n" not in source:
source = source.replace("import logging\n", "import logging\nH3_WRAPPER_CAPTURE_INDEX = 0\n")
old = """ if 'model_function_wrapper' in model_options:
output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
else:
output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
for o in range(batch_chunks):
"""
new = """ if 'model_function_wrapper' in model_options:
output_tensor = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond})
else:
output_tensor = model.apply_model(input_x, timestep_, **c)
capture_dir = os.getenv("H3_WRAPPER_CAPTURE_DIR") or os.getenv("H3_CAPTURE_DIR")
if capture_dir and input_x.ndim == 3:
global H3_WRAPPER_CAPTURE_INDEX
os.makedirs(capture_dir, exist_ok=True)
torch.cuda.synchronize(input_x.device)
torch.save({"input_x": input_x.detach().cpu(), "timestep": timestep_.detach().cpu(), "output": output_tensor.detach().cpu(), "mult": torch.stack([m.detach().cpu() for m in mult]), "cond_or_uncond": cond_or_uncond}, os.path.join(capture_dir, f"wrapper_{H3_WRAPPER_CAPTURE_INDEX:02d}.pt"))
H3_WRAPPER_CAPTURE_INDEX += 1
output = output_tensor.chunk(batch_chunks)
for o in range(batch_chunks):
"""
if source.count(old) == 1:
source = source.replace(old, new)
elif new not in source:
raise RuntimeError("Unable to locate Comfy model wrapper output block.")
path.write_text(source, encoding="utf-8")
print("Applied H3 wrapper boundary capture patch.")