From 1e6a65489a0b8b738f028abef183358ed83be160 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 16:23:08 +0700 Subject: [PATCH] Capture Comfy denoised wrapper boundary --- tools/patch_comfy_wrapper_capture.py | 40 ++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 tools/patch_comfy_wrapper_capture.py diff --git a/tools/patch_comfy_wrapper_capture.py b/tools/patch_comfy_wrapper_capture.py new file mode 100644 index 0000000..a459036 --- /dev/null +++ b/tools/patch_comfy_wrapper_capture.py @@ -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.")