28 lines
1.4 KiB
Python
28 lines
1.4 KiB
Python
|
"""Persist compact per-step sampler state for one direct FL2VA parity run."""
|
||
|
|
|
||
|
|
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")
|
||
|
|
|
||
|
|
old = " if callback is not None:\n callback(x[\"i\"], x[\"denoised\"], x[\"x\"], total_steps)\n"
|
||
|
|
new = (
|
||
|
|
" capture_dir = os.getenv(\"H3_SAMPLER_CAPTURE_DIR\")\n"
|
||
|
|
" if capture_dir:\n"
|
||
|
|
" os.makedirs(capture_dir, exist_ok=True)\n"
|
||
|
|
" if x[\"i\"] == 0:\n"
|
||
|
|
" torch.save({\"sigmas\": sigmas.detach().cpu(), \"initial_x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, \"initial.pt\"))\n"
|
||
|
|
" torch.save({\"sigma\": x[\"sigma\"].detach().cpu(), \"denoised\": x[\"denoised\"].detach().cpu(), \"x\": x[\"x\"].detach().cpu()}, os.path.join(capture_dir, f\"step_{x[\"i\"]:02d}.pt\"))\n"
|
||
|
|
" if callback is not None:\n"
|
||
|
|
" callback(x[\"i\"], x[\"denoised\"], x[\"x\"], total_steps)\n"
|
||
|
|
)
|
||
|
|
if source.count(old) == 1:
|
||
|
|
source = source.replace(old, new)
|
||
|
|
elif new not in source:
|
||
|
|
raise RuntimeError("Unable to locate Comfy sampler callback.")
|
||
|
|
path.write_text(source, encoding="utf-8")
|
||
|
|
print("Applied H3 sampler reference capture patch.")
|