From ebe139e5f5190c9f11bbb92084eb69c6a9542d16 Mon Sep 17 00:00:00 2001 From: Krishi Attri Date: Thu, 11 Jun 2026 01:42:19 +0000 Subject: [PATCH] v0.2.0: GPU-validated inside ComfyUI; copy-on-patch cache safety; single-step reset fix End-to-end validation on RTX 5090 (ComfyUI 0.24.0, Hunyuan3DWrapper master, Hunyuan3D-2mini fp16, 30-step image-to-mesh workflow via the ComfyUI API): - patch engages exactly on schedule (hermite i3: 11 computed + 19 skipped; dmd i5: 7 + 23), sampling 16.1 -> 43.1 it/s (2.68x) for hermite i3 - quality A/B same seed vs unaccelerated mesh: hermite i3 F1@0.05 0.825 (different-seed floor 0.751); dmd i5 0.719, at the floor on mini, so the upstream DMD lossless claim is not repeated for this node - repeat runs in one session reset per-run state correctly Two bugs found by the validation, both fixed with regression tests: 1. ComfyUI cache aliasing: the node mutated and returned the same pipeline object, so a cached node output could silently run a different config after an intervening run re-patched the shared pipeline (observed live: a cached hermite-i3 output ran dmd-i5). apply/remove now return a shallow copy that owns its patch; the input pipeline is never mutated. 2. Run-boundary reset missed back-to-back single-step runs (t=0 then t=0); the boundary check is now non-increasing instead of strictly decreasing. Defaults switch to the best-measured config (hermite, interval 3). Tests 35 -> 38. README validation section rewritten from measurements; REGISTRY-SUBMIT.md documents the registry.comfy.org publish flow. --- .comfyignore | 4 ++ README.md | 112 ++++++++++++++++++++++++++++++++++---------- REGISTRY-SUBMIT.md | 50 ++++++++++++++++++++ __init__.py | 2 +- hicache_patch.py | 79 +++++++++++++++++++++---------- nodes.py | 16 ++++--- pyproject.toml | 4 +- tests/test_nodes.py | 26 +++++----- tests/test_patch.py | 85 +++++++++++++++++++++++++++------ 9 files changed, 291 insertions(+), 87 deletions(-) create mode 100644 .comfyignore create mode 100644 REGISTRY-SUBMIT.md diff --git a/.comfyignore b/.comfyignore new file mode 100644 index 0000000..87f6c87 --- /dev/null +++ b/.comfyignore @@ -0,0 +1,4 @@ +__pycache__/ +tests/ +REGISTRY-SUBMIT.md +.comfyignore diff --git a/README.md b/README.md index 01ba796..0fb0c65 100644 --- a/README.md +++ b/README.md @@ -6,11 +6,13 @@ the [`hicache-pp`](https://pypi.org/project/hicache-pp/) library ([HiCache](https://arxiv.org/abs/2508.16984) Hermite-polynomial and HiCache++ DMD/Prony-exponential forecasters). -> **Status: beta.** The forecasting math, schedule, and patch logic are -> CPU-unit-tested against the library and against the wiring pattern of the -> measured adapter repos (see *Validated on* below). End-to-end inside-ComfyUI -> GPU validation is still pending — until then, treat the node as experimental -> and A/B your own outputs. +> **Status: GPU-validated end-to-end inside ComfyUI** (2026-06-11, RTX 5090, +> real ComfyUI + kijai/ComfyUI-Hunyuan3DWrapper + Hunyuan3D-2mini fp16): +> the patch engages on schedule, sampling runs 2.7x faster at the recommended +> `hermite interval=3`, the accelerated mesh stays well above the +> different-seed noise floor in F-score against the unaccelerated mesh, and +> repeated runs in one session reset state correctly. Measured details in +> *Validated on* below. ## What it does @@ -39,8 +41,8 @@ covers the wrapper's vendored `hy3dgen` (Hunyuan3D 2.0) and `hy3dshape` | method | basis | when to use | |---|---|---| -| `hermite` | HiCache — dual-scaled physicist's Hermite polynomial ([arXiv:2508.16984](https://arxiv.org/abs/2508.16984)) | conservative intervals (i3) | -| `dmd` *(default)* | HiCache++ — Dynamic Mode Decomposition / Prony exponential basis | larger intervals (i4–i5); degrades most gracefully | +| `hermite` *(default)* | HiCache, dual-scaled physicist's Hermite polynomial ([arXiv:2508.16984](https://arxiv.org/abs/2508.16984)) | interval 3; the best-measured setting inside ComfyUI and on both upstream checkpoints | +| `dmd` | HiCache++, Dynamic Mode Decomposition / Prony exponential basis | larger intervals (4 to 5) on the big 2.1 DiT; on mini it measured below the seed-noise floor inside ComfyUI (see *Validated on*) | | `auto` | holdout-selected per compute step: serves DMD only when it demonstrably beats the polynomial on the cached window | when unsure | ### Inputs @@ -48,8 +50,8 @@ covers the wrapper's vendored `hy3dgen` (Hunyuan3D 2.0) and `hy3dshape` | input | default | meaning | |---|---|---| | `pipeline` | — | `HY3DMODEL` from `Hy3DModelLoader` | -| `method` | `dmd` | forecast basis (above) | -| `interval` | 5 | compute 1 step, forecast `interval-1`. `1` disables caching | +| `method` | `hermite` | forecast basis (above) | +| `interval` | 3 | compute 1 step, forecast `interval-1`. `1` disables caching | | `warmup_steps` | 2 | always compute the first N steps (floored at 1 — the first step has nothing to forecast from) | | `enable` | true | off = unpatch, restore the original DiT | | `max_order` | 1 | Hermite / finite-difference order | @@ -60,7 +62,42 @@ Re-running the node with new parameters re-patches cleanly; `enable=false` (or `interval=1`) restores the original model. State resets automatically at each new sampling run. -## Measured numbers (from the adapter repos — not yet re-measured inside ComfyUI) +The node never mutates the pipeline it receives: it returns a shallow copy +whose `model` attribute is the patch (weights stay shared, so this costs no +VRAM). This matters because ComfyUI caches node outputs keyed on node inputs; +an in-place patch lets a cached output alias a pipeline that a later run +re-patched with different settings. That failure was actually observed during +GPU validation (a cached `hermite interval=3` output silently ran +`dmd interval=5`) and is covered by a regression test. + +## Measured inside ComfyUI (this node, end-to-end) + +Setup: ComfyUI 0.24.0, kijai/ComfyUI-Hunyuan3DWrapper master, Hunyuan3D-2mini +fp16 single-file checkpoint, RTX 5090, torch 2.12 cu128. Workflow: LoadImage +(518x518 RGBA crop) -> InvertMask -> Hy3DModelLoader -> HiCacheAccelerate -> +Hy3DGenerateMesh (30 steps, cfg 5.5) -> Hy3DVAEDecode (octree 384) -> +Hy3DExportMesh. Quality metric: F-score@0.05 of the accelerated mesh against +the unaccelerated mesh from the same seed (50k surface samples each, shared +canonical frame, no alignment needed); as a floor, two unaccelerated runs +that differ only in seed score 0.751 on the same metric. + +| config | DiT steps run | sampling | sampling speedup | F1@0.05 vs baseline | +|---|---|---:|---:|---:| +| baseline (`enable=false`) | 30/30 | 1.87 s (16.1 it/s) | 1.00x | 1.000 | +| `hermite interval=3` | 11/30 | 0.70 s (43.1 it/s) | **2.68x** | **0.825** | +| `dmd interval=5` | 7/30 | 1.10 s (27.2 it/s) | 1.69x | 0.719 | + +Reading: `hermite i3` is comfortably above the 0.751 different-seed floor +(the accelerated mesh is closer to its baseline than a fresh seed is), at a +2.7x sampling speedup. `dmd i5` lands at the floor on mini inside ComfyUI; +its forecast also costs more per skipped step (an SVD), which eats into the +speedup on a DiT this small. Hence the node defaults are `hermite` / `3`. +On mini the end-to-end win is modest in absolute terms (about 1.2 s of a +roughly 9 s run) because VAE decode and marching cubes dominate; the larger +the shape DiT and the higher the step count, the larger the end-to-end share +the node saves. + +## Measured numbers (from the adapter repos, outside ComfyUI) The same forecasters, wired into the upstream Hunyuan3D pipelines (outside ComfyUI), measured on Toys4K image-to-3D with geometry-preserving A/B @@ -83,9 +120,12 @@ Source adapter repos (measured, runnable): [hunyuan2-plus-plus](https://github.com/Archerkattri/hunyuan2-plus-plus) (2.0 + mini), [hunyuan2.1-plus-plus](https://github.com/Archerkattri/hunyuan2.1-plus-plus) (2.1). -These numbers transfer to ComfyUI only insofar as the wrapper's vendored -pipelines run the same denoise loop (they do, structurally — that is what the -unit tests pin down), but **no inside-ComfyUI measurement has been done yet**. +These numbers come from first-class pipeline wiring (the forecast caches the +CFG-combined velocity); the ComfyUI node is a model-level patch that caches +the pre-CFG stacked output instead. For Hermite the two are mathematically +identical. For DMD they are not, and the inside-ComfyUI measurement above +confirms the difference is real on mini: prefer `hermite` in this node unless +you re-A/B `dmd` on your own checkpoint. ## Install @@ -107,7 +147,25 @@ to provide the `HY3DMODEL` pipeline the node accelerates. What has actually been verified, and how: -* **CPU unit tests (35, in `tests/`)** — run `pytest` in this repo, no ComfyUI +* **End-to-end GPU validation inside ComfyUI** (2026-06-11, RTX 5090, ComfyUI + 0.24.0, wrapper master, Hunyuan3D-2mini fp16): + * headless ComfyUI boots with the pack installed, zero import errors, and + `HiCacheAccelerate` appears in `/object_info`; + * the full image-to-mesh workflow executes through the ComfyUI API with the + node wired in, and the patch demonstrably engages: with 30 sampling steps + the run-boundary log reports `11 computed + 19 skipped` for + `hermite interval=3` and `7 computed + 23 skipped` for `dmd interval=5`, + exactly the library schedule, with the sampling rates and meshes in the + table above; + * repeated runs in one server session (the same patched pipeline served + from ComfyUI's node-output cache across prompts) each reset the forecast + state at the run boundary and produce fresh, seed-dependent meshes; + * two bugs were found by this validation and are fixed with regression + tests: (1) ComfyUI cache aliasing, fixed by copy-on-patch (see *Inputs*); + (2) back-to-back single-step runs both start at t=0, which the old + strictly-decreasing run-boundary check missed, so the second run could be + served a stale anchor; the boundary check is now non-increasing. +* **CPU unit tests (38, in `tests/`)**, run `pytest` in this repo, no ComfyUI needed: * the node pack loads standalone exactly the way ComfyUI loads it (`NODE_CLASS_MAPPINGS`, `INPUT_TYPES` schema, zero `comfy` imports); @@ -117,17 +175,18 @@ What has actually been verified, and how: are forecast-filled, and the patch's outputs are **bit-identical** to driving `hicache-pp`'s state machine directly (the adapter wiring pattern); * DMD forecasts are exact (rel. err < 1e-3) on exponential velocity - trajectories — its solution class — once its snapshot window fills, and - beat naive last-output reuse by >10× there; - * run-boundary reset, re-patch/unpatch, attribute & `.to()` passthrough, - and parameter validation. -* **Adapter-repo GPU measurements** — the identical forecasters and schedule, + trajectories (its solution class) once its snapshot window fills, and + beat naive last-output reuse by >10x there; + * run-boundary reset (including the single-step edge), copy-on-patch + non-aliasing, re-patch/unpatch, attribute & `.to()` passthrough, and + parameter validation. +* **Adapter-repo GPU measurements**: the identical forecasters and schedule, measured in the upstream (non-ComfyUI) Hunyuan3D pipelines (tables above). -**Not yet validated:** running inside ComfyUI on GPU with the real wrapper and -real checkpoints (including mini/turbo), and quality A/B of the resulting -meshes through the full Comfy graph. Until that lands this repo stays 0.x/beta -and is not submitted to the Comfy Registry. +**Not yet measured:** the full-size Hunyuan3D 2.0/2.1 checkpoints inside +ComfyUI (the mechanism is identical and checkpoint-agnostic, and the adapter +repos measured those checkpoints upstream, but the in-Comfy numbers above are +mini-only), and the multiview sampler path (`Hy3DGenerateMeshMultiView`). ### One honest design note @@ -136,9 +195,10 @@ The adapter repos wire the forecast into the pipeline loop and cache the a ComfyUI node) caches the *pre-CFG stacked* DiT output instead. For the Hermite forecaster the two are mathematically identical (forecast and CFG combine are both linear); for DMD the fit runs on the stacked trajectory -rather than the combined one — same dynamics, but not the literally identical -computation, which is one more reason the GPU A/B is required before any -lossless claim is repeated for this node. +rather than the combined one. The inside-ComfyUI A/B above shows this matters +in practice: DMD's upstream "lossless at interval 5 on mini" result does NOT +transfer to the model-level patch, so this README does not claim it. Hermite +numbers transfer exactly, and that is what the defaults use. ## License diff --git a/REGISTRY-SUBMIT.md b/REGISTRY-SUBMIT.md new file mode 100644 index 0000000..6aecc92 --- /dev/null +++ b/REGISTRY-SUBMIT.md @@ -0,0 +1,50 @@ +# Submitting ComfyUI-HiCache to the Comfy Registry + +Status: the node is GPU-validated end-to-end (see README, *Validated on*) and +version 0.2.0 is ready to publish. Everything below the "human-only" line is +prepared; the two human-only steps need Krishi's browser login. + +## What is already in place + +* `pyproject.toml` has the required `[project]` fields (`name`, + `description`, `version = "0.2.0"`, `license`, `dependencies`) and the + `[tool.comfy]` block: + * `PublisherId = "archerkattri"` + * `DisplayName = "ComfyUI-HiCache"` +* `.comfyignore` excludes tests and tooling files from the published package. +* The repo is public at `https://github.com/Archerkattri/ComfyUI-HiCache`. + +## Human-only steps (browser, one time) + +1. Create the publisher account: go to https://registry.comfy.org, log in, + and create a publisher. The publisher ID is the part after the `@` on the + profile page; it is globally unique and immutable. It MUST be + `archerkattri` to match `pyproject.toml`. If that ID is taken, pick + another and update `PublisherId` in `pyproject.toml` before publishing. +2. Create an API key: registry nodes section -> select the publisher -> + create a new API key. Name it (for example `comfyui-hicache-publish`) and + store it safely; lost keys cannot be recovered, only replaced. + +## Publish (CLI, after the human-only steps) + +```bash +pip install comfy-cli # if not already installed +cd ComfyUI-HiCache # repo root, where pyproject.toml lives +comfy node publish # prompts for the API key +``` + +Expected confirmation: `Version 0.2.0 Published.` The node then appears at +`https://registry.comfy.org/publishers/archerkattri/nodes/comfyui-hicache` +and becomes installable through ComfyUI-Manager once the registry index +picks it up. + +Do NOT use the registry's GitHub Actions auto-publish flow +(`publish_action.yml` + `REGISTRY_ACCESS_TOKEN` secret): this repo +deliberately carries no CI. Publish from the CLI for each release. + +## Releasing future versions + +1. Bump `version` in `pyproject.toml` AND `__version__` in `__init__.py` + (keep them identical, semver). +2. Run the tests: `pytest` (38 must pass). +3. Commit, push, then `comfy node publish` again with the same API key. diff --git a/__init__.py b/__init__.py index d1550b0..01f0cce 100644 --- a/__init__.py +++ b/__init__.py @@ -18,7 +18,7 @@ except ImportError: # imported outside ComfyUI (tests / tooling), not as a pack sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__version__ = "0.1.0" +__version__ = "0.2.0" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] WEB_DIRECTORY = None diff --git a/hicache_patch.py b/hicache_patch.py index 71c6844..bce8a2f 100644 --- a/hicache_patch.py +++ b/hicache_patch.py @@ -40,15 +40,30 @@ fit is done on the stacked trajectory rather than the combined one — in practice the same dynamics; this is the standard trade-off of a model-level patch versus first-class pipeline wiring. -Run-boundary detection: these pipelines sample with monotonically *increasing* -timesteps (``sigmas = linspace(0, 1, N)``; "we start from 0" in the upstream -source). A timestep value lower than the previous call's therefore marks a new -sampling run, and the forecast state is re-initialised. The total step count is -not knowable from inside the model, so the schedule has no end-of-run -always-compute window (``end_enhance``) — only the initial warmup window. +Run-boundary detection: these pipelines sample with *strictly increasing* +timesteps within a run (``sigmas = linspace(0, 1, N)``; "we start from 0" in +the upstream source) and call the model exactly once per step. A timestep +value lower than **or equal to** the previous call's therefore marks a new +sampling run, and the forecast state is re-initialised (the equality case +catches back-to-back single-step runs, which both start at t=0 and would +otherwise be served a stale anchor from the previous run). The total step +count is not knowable from inside the model, so the schedule has no +end-of-run always-compute window (``end_enhance``) — only the initial warmup +window. + +Cache safety inside ComfyUI: :func:`apply_hicache` / :func:`remove_hicache` +never mutate the pipeline they are given — they return a shallow copy whose +``model`` attribute is replaced (weights are shared, so this costs nothing). +ComfyUI caches node *outputs* keyed on node *inputs*; an in-place patch lets +a cached output alias a pipeline that a later run re-patched with different +settings (GPU-validated failure mode: a cached "hermite interval=3" output +silently ran "dmd interval=5" after an intervening run re-patched the shared +object). With copy-on-patch every cached output permanently owns its own +configuration. """ from __future__ import annotations +import copy import logging from typing import Any, Dict, Optional @@ -99,8 +114,8 @@ class HiCacheModelPatch(torch.nn.Module): submodule, so ``.to()`` / ``.half()`` / state-dict access recurse into it). """ - def __init__(self, model: torch.nn.Module, *, method: str = "dmd", - interval: int = 5, warmup_steps: int = 2, max_order: int = 1, + def __init__(self, model: torch.nn.Module, *, method: str = "hermite", + interval: int = 3, warmup_steps: int = 2, max_order: int = 1, sigma: float = 0.5, dmd_history: int = 5) -> None: validate_config(method, interval, warmup_steps, max_order, sigma, dmd_history) super().__init__() @@ -174,9 +189,12 @@ class HiCacheModelPatch(torch.nn.Module): def forward(self, latent_model_input: torch.Tensor, timestep: Any, *args: Any, **kwargs: Any) -> torch.Tensor: t_val = self._timestep_value(timestep) - # New sampling run: timesteps in these pipelines increase monotonically - # (sigma 0 -> 1), so a drop in t means the loop restarted. - if self._state is None or self._last_t is None or t_val < self._last_t: + # New sampling run: within a run these pipelines call the model once + # per step with strictly increasing timesteps (sigma 0 -> 1), so a + # non-increasing t means the loop restarted. `<=` (not `<`) so that + # back-to-back single-step runs (t=0 then t=0 again) also reset + # instead of serving the previous run's anchor. + if self._state is None or self._last_t is None or t_val <= self._last_t: self.reset() self._last_t = t_val @@ -200,36 +218,49 @@ class HiCacheModelPatch(torch.nn.Module): # --------------------------------------------------------------------------- # apply / remove on a pipeline object # --------------------------------------------------------------------------- -def apply_hicache(pipeline: Any, *, method: str = "dmd", interval: int = 5, +def apply_hicache(pipeline: Any, *, method: str = "hermite", interval: int = 3, warmup_steps: int = 2, max_order: int = 1, sigma: float = 0.5, dmd_history: int = 5) -> Any: - """Patch ``pipeline.model`` with a :class:`HiCacheModelPatch` (idempotent). + """Return a shallow copy of ``pipeline`` whose ``model`` is patched. - If the pipeline is already patched the old patch is removed first, so - re-running the node with new parameters reconfigures cleanly. + The input pipeline is NOT mutated (ComfyUI caches node outputs keyed on + node inputs, so a cached output must own its configuration forever — see + the module docstring). Weights are shared between the copy and the + original; only the wrapper object and the ``model`` attribute differ. + If the given pipeline is already patched, the patch is replaced (never + nested), so re-running the node with new parameters reconfigures cleanly. """ if not hasattr(pipeline, "model"): raise TypeError( "HiCache: pipeline has no `.model` attribute - expected a Hunyuan3D " f"shape pipeline, got {type(pipeline).__name__}" ) - remove_hicache(pipeline) - patch = HiCacheModelPatch( - pipeline.model, method=method, interval=interval, warmup_steps=warmup_steps, + inner = pipeline.model + if getattr(inner, "_hicache_is_patch", False): + inner = inner.inner # replace, never nest + patched = copy.copy(pipeline) + patched.model = HiCacheModelPatch( + inner, method=method, interval=interval, warmup_steps=warmup_steps, max_order=max_order, sigma=sigma, dmd_history=dmd_history, ) - pipeline.model = patch logger.info( - "[HiCache] patched %s.model: method=%s interval=%d warmup=%d", + "[HiCache] patched copy of %s: method=%s interval=%d warmup=%d", type(pipeline).__name__, method, interval, warmup_steps, ) - return pipeline + return patched def remove_hicache(pipeline: Any) -> Any: - """Restore the original DiT on ``pipeline.model`` (no-op if not patched).""" + """Return ``pipeline`` with the original DiT on ``.model``. + + If the pipeline is patched, a shallow copy with the patch unwrapped is + returned (the input is not mutated); if it is not patched, the pipeline + is returned unchanged. + """ model = getattr(pipeline, "model", None) if model is not None and getattr(model, "_hicache_is_patch", False): - pipeline.model = model.inner - logger.info("[HiCache] removed patch from %s.model", type(pipeline).__name__) + clean = copy.copy(pipeline) + clean.model = model.inner + logger.info("[HiCache] removed patch from %s", type(pipeline).__name__) + return clean return pipeline diff --git a/nodes.py b/nodes.py index 8738b53..e68669c 100644 --- a/nodes.py +++ b/nodes.py @@ -32,18 +32,20 @@ class HiCacheAccelerate: "tooltip": "Hunyuan3D shape pipeline from Hy3DModelLoader " "(kijai/ComfyUI-Hunyuan3DWrapper)."}), "method": (list(METHODS), { - "default": "dmd", + "default": "hermite", "tooltip": "Forecast basis on skipped steps:\n" - "hermite = HiCache (polynomial, arXiv:2508.16984)\n" - "dmd = HiCache++ (exponential / Prony-DMD; degrades " - "most gracefully at large intervals)\n" + "hermite = HiCache (polynomial, arXiv:2508.16984; " + "best measured quality inside ComfyUI)\n" + "dmd = HiCache++ (exponential / Prony-DMD; worth " + "trying at large intervals on the big 2.1 DiT)\n" "auto = holdout-pick the better of the two per " "compute step"}), "interval": ("INT", { - "default": 5, "min": 1, "max": 12, "step": 1, + "default": 3, "min": 1, "max": 12, "step": 1, "tooltip": "Compute one DiT step, then forecast interval-1 " "steps. 1 = caching disabled (every step computed). " - "Measured sweet spots: hermite i3, dmd i4-i5."}), + "Measured sweet spot inside ComfyUI: hermite " + "interval 3."}), "warmup_steps": ("INT", { "default": 2, "min": 0, "max": 100, "step": 1, "tooltip": "Always compute the first N sampling steps before " @@ -78,7 +80,7 @@ class HiCacheAccelerate: "Hermite / HiCache++ DMD, via the hicache-pp library)." ) - def patch(self, pipeline, method="dmd", interval=5, warmup_steps=2, + def patch(self, pipeline, method="hermite", interval=3, warmup_steps=2, enable=True, max_order=1, sigma=0.5, dmd_history=5): if not enable or interval <= 1: return (remove_hicache(pipeline),) diff --git a/pyproject.toml b/pyproject.toml index 02008b6..cb86e07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,11 +5,11 @@ [project] name = "comfyui-hicache" description = "Training-free Hunyuan3D shape-generation acceleration: skip DiT steps and forecast the flow-matching velocity instead (HiCache Hermite / HiCache++ DMD, via hicache-pp). Pairs with kijai/ComfyUI-Hunyuan3DWrapper." -version = "0.1.0" +version = "0.2.0" license = { file = "LICENSE" } requires-python = ">=3.9" dependencies = ["hicache-pp>=1.1.0"] -classifiers = ["Development Status :: 4 - Beta"] +classifiers = ["Development Status :: 5 - Production/Stable"] [project.urls] Repository = "https://github.com/Archerkattri/ComfyUI-HiCache" diff --git a/tests/test_nodes.py b/tests/test_nodes.py index c1619a8..7e2d6ed 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -78,23 +78,25 @@ def test_node_function_patches_and_unpatches(): original = pipe.model (out,) = node.patch(pipe, method="hermite", interval=4, warmup_steps=2) - assert out is pipe - assert getattr(pipe.model, "_hicache_is_patch", False) - assert pipe.model.inner is original + assert out is not pipe, "node must return a copy, not mutate (ComfyUI caching)" + assert pipe.model is original, "input pipeline must stay unpatched" + assert getattr(out.model, "_hicache_is_patch", False) + assert out.model.inner is original # re-patch with new params replaces (not nests) the patch - (out,) = node.patch(pipe, method="dmd", interval=3, warmup_steps=1) - assert pipe.model.inner is original - assert pipe.model.method == "dmd" and pipe.model.interval == 3 + (out2,) = node.patch(out, method="dmd", interval=3, warmup_steps=1) + assert out2.model.inner is original + assert out2.model.method == "dmd" and out2.model.interval == 3 + assert out.model.method == "hermite", "earlier output keeps its config" # enable=False restores the original model - (out,) = node.patch(pipe, enable=False) - assert pipe.model is original + (out3,) = node.patch(out2, enable=False) + assert out3.model is original - # interval=1 means "no caching" -> also unpatched - node.patch(pipe, interval=4) - (out,) = node.patch(pipe, interval=1) - assert pipe.model is original + # interval=1 means "no caching" -> also unpatched (and an unpatched + # input passes through unchanged) + (out4,) = node.patch(pipe, interval=1) + assert out4 is pipe and out4.model is original def test_node_rejects_bad_pipeline(): diff --git a/tests/test_patch.py b/tests/test_patch.py index d7ed7f7..7f74bcc 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -113,7 +113,7 @@ def test_compute_calls_follow_schedule(method, interval, warmup): steps = 30 dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method=method, interval=interval, warmup_steps=warmup) + pipe = hp.apply_hicache(pipe, method=method, interval=interval, warmup_steps=warmup) pipe(num_inference_steps=steps) @@ -155,7 +155,7 @@ def _run_pair(ref_dit, acc_dit, method, steps=30, interval=4, warmup=2, **kw): torch.manual_seed(1) _, ref_outputs = MockPipeline(ref_dit)(num_inference_steps=steps) pipe = MockPipeline(acc_dit) - hp.apply_hicache(pipe, method=method, interval=interval, warmup_steps=warmup, **kw) + pipe = hp.apply_hicache(pipe, method=method, interval=interval, warmup_steps=warmup, **kw) torch.manual_seed(1) _, acc_outputs = pipe(num_inference_steps=steps) computes = expected_compute_steps(steps, interval, warmup) @@ -192,7 +192,7 @@ def test_patch_reproduces_hicache_pp_forecast_bitwise(): dit_b.load_state_dict(dit_a.state_dict()) pipe = MockPipeline(dit_a) - hp.apply_hicache(pipe, method="hermite", interval=interval, warmup_steps=warmup) + pipe = hp.apply_hicache(pipe, method="hermite", interval=interval, warmup_steps=warmup) torch.manual_seed(1) _, acc_outputs = pipe(num_inference_steps=steps) @@ -272,7 +272,7 @@ def test_skipped_outputs_finite_and_better_than_naive_reuse(): def test_interval_1_or_warmup_dominates_means_no_skips(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="hermite", interval=2, warmup_steps=100) + pipe = hp.apply_hicache(pipe, method="hermite", interval=2, warmup_steps=100) pipe(num_inference_steps=20) assert dit.calls == 20 and pipe.model.skipped_steps == 0 @@ -283,7 +283,7 @@ def test_interval_1_or_warmup_dominates_means_no_skips(): def test_dmd_method_populates_snapshots_and_forecasts(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2, dmd_history=5) + pipe = hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2, dmd_history=5) pipe(num_inference_steps=40) st = pipe.model._state assert st["backend"] == "dmd" @@ -294,7 +294,7 @@ def test_dmd_method_populates_snapshots_and_forecasts(): def test_hermite_method_keeps_snapshots_empty(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="hermite", interval=4, warmup_steps=2) + pipe = hp.apply_hicache(pipe, method="hermite", interval=4, warmup_steps=2) pipe(num_inference_steps=40) st = pipe.model._state assert st["backend"] == "hermite" @@ -305,7 +305,7 @@ def test_hermite_method_keeps_snapshots_empty(): def test_auto_method_makes_a_holdout_choice(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="auto", interval=3, warmup_steps=2, dmd_history=8) + pipe = hp.apply_hicache(pipe, method="auto", interval=3, warmup_steps=2, dmd_history=8) pipe(num_inference_steps=40) st = pipe.model._state assert st["backend"] == "auto" @@ -319,7 +319,7 @@ def test_auto_method_makes_a_holdout_choice(): def test_state_resets_between_sampling_runs(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2) + pipe = hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2) pipe(num_inference_steps=20) first_run_calls = dit.calls @@ -334,7 +334,7 @@ def test_state_resets_between_sampling_runs(): def test_attribute_and_device_passthrough(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="hermite", interval=4, warmup_steps=2) + pipe = hp.apply_hicache(pipe, method="hermite", interval=4, warmup_steps=2) # pipelines check hasattr(self.model, 'guidance_embed') and call .to() assert pipe.model.guidance_embed is False pipe.model.to(torch.float32) @@ -344,12 +344,67 @@ def test_attribute_and_device_passthrough(): def test_remove_restores_original_model(): dit = MockDiT() pipe = MockPipeline(dit) - hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2) - assert pipe.model is not dit - hp.remove_hicache(pipe) - assert pipe.model is dit - hp.remove_hicache(pipe) # idempotent - assert pipe.model is dit + patched = hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2) + assert patched is not pipe and patched.model is not dit + assert pipe.model is dit, "apply must not mutate its input pipeline" + restored = hp.remove_hicache(patched) + assert restored.model is dit + assert patched.model is not dit, "remove must not mutate its input either" + assert hp.remove_hicache(pipe) is pipe # unpatched input passes through + + +def test_apply_does_not_alias_between_configs(): + """ComfyUI cache-aliasing regression (found in GPU validation): node + outputs are cached keyed on node inputs, so a pipeline returned for + config A must keep running config A even after the same loader pipeline + is later patched with config B. In-place patching broke this: a cached + 'hermite interval=3' output silently ran 'dmd interval=5'.""" + torch.manual_seed(0) + dit = MockDiT() + base = MockPipeline(dit) + + pipe_a = hp.apply_hicache(base, method="hermite", interval=3, warmup_steps=2) + pipe_b = hp.apply_hicache(base, method="dmd", interval=5, warmup_steps=2) + + # the loader's pipeline is untouched; each copy owns its config + assert base.model is dit + assert (pipe_a.model.method, pipe_a.model.interval) == ("hermite", 3) + assert (pipe_b.model.method, pipe_b.model.interval) == ("dmd", 5) + assert pipe_a.model.inner is dit and pipe_b.model.inner is dit + + # serving the A copy after B was created still runs A's schedule + steps = 30 + pipe_a(num_inference_steps=steps) + assert dit.calls == len(expected_compute_steps(steps, 3, 2)) + + +def test_apply_on_patched_pipeline_replaces_not_nests(): + dit = MockDiT() + base = MockPipeline(dit) + once = hp.apply_hicache(base, method="hermite", interval=3, warmup_steps=2) + twice = hp.apply_hicache(once, method="dmd", interval=5, warmup_steps=2) + assert twice.model.inner is dit, "patch must replace, never nest" + assert (once.model.method, once.model.interval) == ("hermite", 3) + + +def test_single_step_runs_reset_at_equal_timestep(): + """Run-boundary regression: two back-to-back single-step runs both call + the model at t=0. The boundary check must treat the *equal* timestep as + a new run and recompute — not serve the previous run's cached anchor.""" + dit = MockDiT() + patch = hp.HiCacheModelPatch(dit, method="hermite", interval=3, + warmup_steps=1) + lmi = torch.randn(2, 8, 4) + t0 = torch.zeros(2) + + out1 = patch(lmi, t0, cond=None) # run 1 (single step) + with torch.no_grad(): + dit.base.add_(1.0) # "new image": DiT output changes + out2 = patch(lmi, t0, cond=None) # run 2 (single step, same t) + + assert dit.calls == 2, "second single-step run must recompute, not forecast" + assert not torch.allclose(out1, out2), "stale anchor served across runs" + assert patch.computed_steps == 1 and patch.skipped_steps == 0 # per-run stats # ---------------------------------------------------------------------------