40 lines
2.2 KiB
Python
40 lines
2.2 KiB
Python
"""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.")
|