commit 26615c4bcf3b2186152c07819594ee81108813cd Author: Krishi Attri Date: Wed Jun 10 04:20:21 2026 +0000 ComfyUI-HiCache v0.1.0 (beta): training-free Hunyuan3D acceleration node One node, HiCacheAccelerate: patches the HY3DMODEL shape pipeline's DiT (kijai/ComfyUI-Hunyuan3DWrapper) to skip DiT forwards and forecast the flow-matching velocity via hicache-pp (HiCache Hermite / HiCache++ DMD / auto holdout selection). - zero comfy-internal imports: the patch targets the pipeline object the wrapper passes between nodes, so it covers hy3dgen (2.0), hy3dshape (2.1) and mini/turbo checkpoints alike - 35 CPU unit tests: standalone node-pack load contract, schedule adherence on a mock DiT with the wrapper's exact loop shape, bit-identity with driving hicache-pp directly, DMD exactness on exponential trajectories, run-boundary reset, unpatch, parameter validation - registry.comfy.org pyproject metadata (publish deferred until GPU-in- ComfyUI validation; repo ships as beta) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2ec1dba --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +.pytest_cache/ +*.egg-info/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..2dd4290 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Krishi Attri + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..01ba796 --- /dev/null +++ b/README.md @@ -0,0 +1,145 @@ +# ComfyUI-HiCache + +**Training-free Hunyuan3D shape-generation acceleration for ComfyUI** — skip DiT +forwards during flow-matching sampling and forecast the velocity instead, using +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. + +## What it does + +One node, `HiCache Accelerate (Hunyuan3D)`, wired between +[kijai/ComfyUI-Hunyuan3DWrapper](https://github.com/kijai/ComfyUI-Hunyuan3DWrapper)'s +`Hy3DModelLoader` and its mesh sampler (`Hy3DGenerateMesh` / +`Hy3DGenerateMeshMultiView`): + +``` +Hy3DModelLoader ──pipeline──▶ HiCache Accelerate ──pipeline──▶ Hy3DGenerateMesh +``` + +The node patches `pipeline.model` (the shape DiT). On *compute* steps the DiT +runs normally and its output is cached as a forecast anchor; on *skipped* steps +the DiT is **not called at all** — the flow-matching velocity is forecast from +the cached anchors. With `interval = N`, roughly `(N-1)/N` of the DiT forwards +are skipped. The patch needs nothing from ComfyUI internals and works on any +pipeline object whose denoise loop calls +`self.model(latent_model_input, timestep, cond, ...)` once per step — which +covers the wrapper's vendored `hy3dgen` (Hunyuan3D 2.0) and `hy3dshape` +(Hunyuan3D 2.1) pipelines, including the mini/turbo checkpoints +([wrapper issue #97](https://github.com/kijai/ComfyUI-Hunyuan3DWrapper/issues/97) +— the patch is checkpoint-agnostic; only the loop shape matters). + +### Methods + +| 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 | +| `auto` | holdout-selected per compute step: serves DMD only when it demonstrably beats the polynomial on the cached window | when unsure | + +### Inputs + +| input | default | meaning | +|---|---|---| +| `pipeline` | — | `HY3DMODEL` from `Hy3DModelLoader` | +| `method` | `dmd` | forecast basis (above) | +| `interval` | 5 | 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 | +| `sigma` | 0.5 | Hermite contraction factor in (0,1) | +| `dmd_history` | 5 | DMD snapshot window length | + +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 same forecasters, wired into the upstream Hunyuan3D pipelines (outside +ComfyUI), measured on Toys4K image-to-3D with geometry-preserving A/B +(F-score@0.05 vs. the uncached baseline, solo re-timed speedups). Full tables +and methodology: +[hicache-plus-plus results](https://github.com/Archerkattri/hicache-plus-plus/blob/master/results/RESULTS.md). + +**Hunyuan3D-2.1** (baseline F-score 0.911): + +| interval | Hermite (HiCache) | DMD (HiCache++) | speedup | +|---:|---:|---:|---:| +| i3 | **0.876** | 0.852 | 1.72× | +| i4 | 0.776 | **0.827** | 1.80× | +| i5 | 0.735 | **0.860** | 1.79× | + +**Hunyuan3D-2-mini** (baseline F-score 0.794, 1.89 s/gen): DMD i5 is **exactly +lossless** (0.794) at 1.69 s; HiCache i3 0.792 at 1.58 s. + +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**. + +## Install + +**Via ComfyUI-Manager:** search for `ComfyUI-HiCache` (once indexed), or +*Install via Git URL* → `https://github.com/Archerkattri/ComfyUI-HiCache`. + +**Manual:** + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/Archerkattri/ComfyUI-HiCache +pip install -r ComfyUI-HiCache/requirements.txt # just: hicache-pp +``` + +Requires [ComfyUI-Hunyuan3DWrapper](https://github.com/kijai/ComfyUI-Hunyuan3DWrapper) +to provide the `HY3DMODEL` pipeline the node accelerates. + +## Validated on + +What has actually been verified, and how: + +* **CPU unit tests (35, 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); + * against a mock DiT with the wrapper's exact call signature, driven by a + synthetic denoise loop replicating the wrapper's flow-matching loop: + DiT forwards happen **only** on the scheduled compute steps, skipped steps + 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, + 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. + +### One honest design note + +The adapter repos wire the forecast into the pipeline loop and cache the +*CFG-combined* velocity; a model-level patch (the only non-invasive option for +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. + +## License + +MIT. Not affiliated with Tencent (Hunyuan3D) or kijai (ComfyUI-Hunyuan3DWrapper). diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d1550b0 --- /dev/null +++ b/__init__.py @@ -0,0 +1,24 @@ +"""ComfyUI-HiCache — training-free Hunyuan3D acceleration for ComfyUI. + +Skips DiT forwards during flow-matching sampling and forecasts the velocity +instead, using the HiCache (Hermite polynomial, arXiv:2508.16984) and +HiCache++ (DMD / Prony exponential) forecasters from the ``hicache-pp`` +library. + +Repository: https://github.com/Archerkattri/ComfyUI-HiCache +License: MIT +""" + +try: + from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +except ImportError: # imported outside ComfyUI (tests / tooling), not as a package + import os + import sys + + 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" +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +WEB_DIRECTORY = None diff --git a/hicache_patch.py b/hicache_patch.py new file mode 100644 index 0000000..71c6844 --- /dev/null +++ b/hicache_patch.py @@ -0,0 +1,235 @@ +"""HiCache / HiCache++ model-forward patch — pure Python, no ComfyUI imports. + +This module contains all of the actual acceleration logic so it can be unit +tested standalone (no ComfyUI install needed). ``nodes.py`` only does the +ComfyUI plumbing around :func:`apply_hicache` / :func:`remove_hicache`. + +How it works +------------ +The Hunyuan3D shape pipelines (both the 2.0 ``hy3dgen`` and 2.1 ``hy3dshape`` +variants vendored by kijai/ComfyUI-Hunyuan3DWrapper, and Tencent's upstream +``Hunyuan3DDiTFlowMatchingPipeline``) run a flow-matching denoise loop of the +form:: + + for t in timesteps: # sigma goes 0 -> 1 + latent_model_input = cat([latents] * 2) # batched CFG + noise_pred = self.model(latent_model_input, timestep, cond, ...) + ...chunk(2) + CFG combine + scheduler.step... + +``self.model`` is a plain instance attribute (the DiT ``nn.Module``), so +replacing ``pipeline.model`` with :class:`HiCacheModelPatch` intercepts every +DiT forward. On *compute* steps the wrapped DiT runs normally and its raw +output (the stacked cond+uncond velocity) is cached as a forecast anchor; on +*skipped* steps the DiT is **not called at all** — the output is forecast from +the cached anchors with ``hicache-pp``: + +* ``hermite`` — HiCache (dual-scaled physicist's Hermite polynomial, + arXiv:2508.16984). +* ``dmd`` — HiCache++ (Dynamic Mode Decomposition / Prony exponential + basis; exact on the feature-ODE solution class, so it stays lossless at + larger skip intervals than the polynomial). +* ``auto`` — holdout-selected per compute step: serve DMD only when it + demonstrably beats the polynomial on the cached window. + +Note on what is cached: the adapter repos (e.g. hunyuan2-plus-plus) wire the +forecast into the pipeline loop and cache the *CFG-combined* velocity. A +model-level patch necessarily caches the *pre-CFG stacked* output instead. +For the Hermite forecaster the two are mathematically identical (the forecast +is linear in the cached features, and CFG combination is linear). For DMD the +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. +""" +from __future__ import annotations + +import logging +from typing import Any, Dict, Optional + +import torch + +from hicache_pp import ( + hicache_init, + hicache_decide, + hicache_update_derivatives, + hicache_forecast, + dmd_update_snapshots, + dmd_forecast_state, + auto_forecast_state, +) + +logger = logging.getLogger("ComfyUI-HiCache") + +METHODS = ("hermite", "dmd", "auto") + +# Sentinel for "unknown total step count": disables the end-of-run +# always-compute window (any real run is far shorter than this). +_NO_END_WINDOW = 1_000_000 + + +def validate_config(method: str, interval: int, warmup_steps: int, + max_order: int, sigma: float, dmd_history: int) -> None: + """Raise ValueError on bad node parameters (mirrors hicache_pp's checks).""" + if method not in METHODS: + raise ValueError(f"method must be one of {METHODS}, got {method!r}") + if interval < 1: + raise ValueError(f"interval must be >= 1, got {interval}") + if warmup_steps < 0: + raise ValueError(f"warmup_steps must be >= 0, got {warmup_steps}") + if max_order < 1: + raise ValueError(f"max_order must be >= 1, got {max_order}") + if not (0.0 < sigma < 1.0): + raise ValueError(f"sigma must be in (0, 1), got {sigma}") + if dmd_history < 3: + raise ValueError(f"dmd_history must be >= 3, got {dmd_history}") + + +class HiCacheModelPatch(torch.nn.Module): + """Drop-in replacement for ``pipeline.model`` that skips DiT forwards. + + Wraps the original DiT module; forwards unknown attribute lookups to it so + pipeline code like ``hasattr(self.model, 'guidance_embed')`` and + ``self.model.to(device)`` keep working (the inner model is a registered + 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, + sigma: float = 0.5, dmd_history: int = 5) -> None: + validate_config(method, interval, warmup_steps, max_order, sigma, dmd_history) + super().__init__() + self.inner = model + self._hicache_is_patch = True # marker for apply/remove + self.method = method + self.interval = int(interval) + self.warmup_steps = int(warmup_steps) + self.max_order = int(max_order) + self.sigma = float(sigma) + self.dmd_history = int(dmd_history) + + self._state: Optional[Dict[str, Any]] = None + self._last_t: Optional[float] = None + # per-run stats (read by the node / logged at run boundaries) + self.computed_steps = 0 + self.skipped_steps = 0 + + # -- attribute passthrough ------------------------------------------------- + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + # fall through to the wrapped DiT (guidance_embed, config, dtype, ...) + inner = super().__getattr__("inner") + return getattr(inner, name) + + # -- state handling ---------------------------------------------------------- + def _fresh_state(self) -> Dict[str, Any]: + # hicache_pp validates interval/order/sigma/backend again here. + return hicache_init( + num_steps=_NO_END_WINDOW, + interval=self.interval, + max_order=self.max_order, + # the very first step of a run must always compute (there is no + # anchor to forecast from), so the effective warmup floor is 1 + first_enhance=max(1, self.warmup_steps), + end_enhance=_NO_END_WINDOW, + sigma=self.sigma, + backend=self.method, + history=self.dmd_history, + ) + + def reset(self) -> None: + """Drop all cached anchors and per-run stats (new sampling run).""" + if self._state is not None and (self.computed_steps or self.skipped_steps): + logger.info( + "[HiCache] run finished: %d computed + %d skipped DiT steps " + "(method=%s, interval=%d)", + self.computed_steps, self.skipped_steps, self.method, self.interval, + ) + self._state = self._fresh_state() + self._last_t = None + self.computed_steps = 0 + self.skipped_steps = 0 + + @staticmethod + def _timestep_value(timestep: Any) -> float: + if torch.is_tensor(timestep): + return float(timestep.reshape(-1)[0].item()) + return float(timestep) + + def _forecast(self, state: Dict[str, Any]) -> torch.Tensor: + if self.method == "dmd": + return dmd_forecast_state(state) + if self.method == "auto": + return auto_forecast_state(state) + return hicache_forecast(state) + + # -- the patched forward ------------------------------------------------------ + 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: + self.reset() + self._last_t = t_val + + state = self._state + if hicache_decide(state) == "forecast": + out = self._forecast(state) + state["step"] += 1 + self.skipped_steps += 1 + return out + + out = self.inner(latent_model_input, timestep, *args, **kwargs) + anchor = out.detach() + hicache_update_derivatives(state, anchor) + if self.method in ("dmd", "auto"): + dmd_update_snapshots(state, anchor, state["history"]) + state["step"] += 1 + self.computed_steps += 1 + return out + + +# --------------------------------------------------------------------------- +# apply / remove on a pipeline object +# --------------------------------------------------------------------------- +def apply_hicache(pipeline: Any, *, method: str = "dmd", interval: int = 5, + 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). + + If the pipeline is already patched the old patch is removed first, 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, + 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", + type(pipeline).__name__, method, interval, warmup_steps, + ) + return pipeline + + +def remove_hicache(pipeline: Any) -> Any: + """Restore the original DiT on ``pipeline.model`` (no-op if not patched).""" + 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__) + return pipeline diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..8738b53 --- /dev/null +++ b/nodes.py @@ -0,0 +1,102 @@ +"""ComfyUI-HiCache node definitions. + +One node, no ComfyUI-internal imports: the patch operates directly on the +pipeline object that kijai/ComfyUI-Hunyuan3DWrapper passes between nodes as +the custom ``HY3DMODEL`` type, so this file imports cleanly outside ComfyUI +too (which is how the unit tests run it). +""" +from __future__ import annotations + +try: + from .hicache_patch import METHODS, apply_hicache, remove_hicache +except ImportError: # imported outside ComfyUI (tests / tooling), not as a package + from hicache_patch import METHODS, apply_hicache, remove_hicache + + +class HiCacheAccelerate: + """Training-free Hunyuan3D shape-DiT acceleration via velocity forecasting. + + Wire between the Hunyuan3DWrapper model loader (``Hy3DModelLoader``) and + the mesh sampler (``Hy3DGenerateMesh`` / ``Hy3DGenerateMeshMultiView``). + On skipped sampling steps the DiT forward is replaced by a cheap forecast + from cached velocities (HiCache Hermite / HiCache++ DMD), cutting the + diffusion-sampling wall clock by roughly ``(interval-1)/interval`` of the + DiT time. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "pipeline": ("HY3DMODEL", { + "tooltip": "Hunyuan3D shape pipeline from Hy3DModelLoader " + "(kijai/ComfyUI-Hunyuan3DWrapper)."}), + "method": (list(METHODS), { + "default": "dmd", + "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" + "auto = holdout-pick the better of the two per " + "compute step"}), + "interval": ("INT", { + "default": 5, "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."}), + "warmup_steps": ("INT", { + "default": 2, "min": 0, "max": 100, "step": 1, + "tooltip": "Always compute the first N sampling steps before " + "any forecasting starts."}), + }, + "optional": { + "enable": ("BOOLEAN", { + "default": True, + "tooltip": "Off = remove the patch and restore the original " + "DiT forward."}), + "max_order": ("INT", { + "default": 1, "min": 1, "max": 4, "step": 1, + "tooltip": "Highest Hermite / finite-difference order " + "(hermite method and warm-up fallback)."}), + "sigma": ("FLOAT", { + "default": 0.5, "min": 0.05, "max": 0.95, "step": 0.05, + "tooltip": "Hermite contraction factor in (0,1); keeps " + "high-order terms bounded."}), + "dmd_history": ("INT", { + "default": 5, "min": 3, "max": 16, "step": 1, + "tooltip": "DMD snapshot window length (dmd/auto methods)."}), + }, + } + + RETURN_TYPES = ("HY3DMODEL",) + RETURN_NAMES = ("pipeline",) + FUNCTION = "patch" + CATEGORY = "HiCache" + DESCRIPTION = ( + "Training-free acceleration of Hunyuan3D shape generation: skip DiT " + "forwards and forecast the flow-matching velocity instead (HiCache " + "Hermite / HiCache++ DMD, via the hicache-pp library)." + ) + + def patch(self, pipeline, method="dmd", interval=5, warmup_steps=2, + enable=True, max_order=1, sigma=0.5, dmd_history=5): + if not enable or interval <= 1: + return (remove_hicache(pipeline),) + return (apply_hicache( + pipeline, + method=method, + interval=interval, + warmup_steps=warmup_steps, + max_order=max_order, + sigma=sigma, + dmd_history=dmd_history, + ),) + + +NODE_CLASS_MAPPINGS = { + "HiCacheAccelerate": HiCacheAccelerate, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "HiCacheAccelerate": "HiCache Accelerate (Hunyuan3D)", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..02008b6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,25 @@ +# Packaging metadata for the registry.comfy.org Comfy Registry. +# Field layout follows https://docs.comfy.org/registry/specifications +# (and the kijai/ComfyUI-Hunyuan3DWrapper precedent). + +[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" +license = { file = "LICENSE" } +requires-python = ">=3.9" +dependencies = ["hicache-pp>=1.1.0"] +classifiers = ["Development Status :: 4 - Beta"] + +[project.urls] +Repository = "https://github.com/Archerkattri/ComfyUI-HiCache" +BugTracker = "https://github.com/Archerkattri/ComfyUI-HiCache/issues" + +[tool.comfy] +PublisherId = "archerkattri" +DisplayName = "ComfyUI-HiCache" +Icon = "" + +[tool.pytest.ini_options] +addopts = "--import-mode=importlib" +testpaths = ["tests"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..99ba9c0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +# torch and ComfyUI are assumed pre-installed (this is a ComfyUI node pack). +hicache-pp>=1.1.0 diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..c1619a8 --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,104 @@ +"""Node-pack registration tests — run standalone, no ComfyUI install. + +ComfyUI loads a custom node pack by importing the package and reading +``NODE_CLASS_MAPPINGS`` / ``NODE_DISPLAY_NAME_MAPPINGS``; these tests verify +that contract plus the INPUT_TYPES schema shape ComfyUI expects. +""" +import importlib.util +import pathlib +import sys + +import pytest + +PACK_DIR = pathlib.Path(__file__).resolve().parents[1] + + +def _load_pack(): + """Import the node pack the way ComfyUI does (as a package by path).""" + name = "comfyui_hicache_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location( + name, PACK_DIR / "__init__.py", + submodule_search_locations=[str(PACK_DIR)], + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod + spec.loader.exec_module(mod) + return mod + + +def test_pack_imports_without_comfy(): + assert "comfy" not in sys.modules, "test must run without ComfyUI" + pack = _load_pack() + assert hasattr(pack, "NODE_CLASS_MAPPINGS") + assert hasattr(pack, "NODE_DISPLAY_NAME_MAPPINGS") + assert "comfy" not in sys.modules, "node pack must not import comfy internals" + + +def test_node_mappings_consistent(): + pack = _load_pack() + assert "HiCacheAccelerate" in pack.NODE_CLASS_MAPPINGS + for key in pack.NODE_CLASS_MAPPINGS: + assert key in pack.NODE_DISPLAY_NAME_MAPPINGS + + +def test_input_types_schema(): + pack = _load_pack() + cls = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"] + schema = cls.INPUT_TYPES() + req = schema["required"] + + # model input + the contract from the task spec: method/interval/warmup + assert req["pipeline"][0] == "HY3DMODEL" + assert set(req["method"][0]) == {"hermite", "dmd", "auto"} + assert req["method"][1]["default"] in req["method"][0] + assert req["interval"][0] == "INT" + assert req["interval"][1]["min"] >= 1 + assert req["warmup_steps"][0] == "INT" + assert req["warmup_steps"][1]["min"] == 0 + assert "enable" in schema["optional"] + + # ComfyUI node contract + assert cls.RETURN_TYPES == ("HY3DMODEL",) + assert callable(getattr(cls, cls.FUNCTION)) + assert isinstance(cls.CATEGORY, str) + + +def test_node_function_patches_and_unpatches(): + import torch + + pack = _load_pack() + node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]() + + class FakePipe: + model = torch.nn.Linear(2, 2) + + pipe = FakePipe() + 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 + + # 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 + + # enable=False restores the original model + (out,) = node.patch(pipe, enable=False) + assert pipe.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 + + +def test_node_rejects_bad_pipeline(): + pack = _load_pack() + node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]() + with pytest.raises(TypeError, match="no `.model` attribute"): + node.patch(object()) diff --git a/tests/test_patch.py b/tests/test_patch.py new file mode 100644 index 0000000..d7ed7f7 --- /dev/null +++ b/tests/test_patch.py @@ -0,0 +1,377 @@ +"""Patch-logic tests on a mock DiT + a synthetic denoise loop (CPU only). + +The mock pipeline reproduces the exact call shape of the Hunyuan3D wrapper's +flow-matching loop (kijai/ComfyUI-Hunyuan3DWrapper, hy3dgen/shapegen and +hy3dshape pipelines): + + latent_model_input = cat([latents] * 2) # batched CFG + timestep = t.expand(2B) / num_train_timesteps # increasing, starts at 0 + noise_pred = self.model(latent_model_input, timestep, cond, guidance=None) + cond_pred, uncond_pred = noise_pred.chunk(2) + noise_pred = uncond_pred + scale * (cond_pred - uncond_pred) + +These tests assert that with the HiCacheModelPatch installed: + * the inner DiT is called exactly on the scheduled compute steps, + * skipped steps are filled by the forecaster (finite, accurate on smooth + velocity trajectories), + * the per-method forecast path (hermite / dmd / auto) is actually exercised, + * state resets across sampling runs, + * parameter plumbing and validation behave. +""" +import importlib.util +import math +import pathlib +import sys + +import pytest +import torch + +PACK_DIR = pathlib.Path(__file__).resolve().parents[1] + + +def _load(modname): + spec = importlib.util.spec_from_file_location( + f"hicache_patch_under_test_{modname}", PACK_DIR / f"{modname}.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +hp = _load("hicache_patch") + + +class MockDiT(torch.nn.Module): + """Stands in for the Hunyuan3D DiT: same call signature as the wrapper's + ``self.model(latent_model_input, timestep, cond, guidance=...)``; returns a + smooth, timestep-dependent velocity field and counts real forwards.""" + + guidance_embed = False # the pipelines hasattr/check this on .model + + def __init__(self, shape=(2, 8, 4)): + super().__init__() + self.shape = shape + self.calls = 0 + self.call_timesteps = [] + base = torch.randn(shape) + self.register_buffer("base", base) + + def forward(self, latent_model_input, timestep, cond, guidance=None): + self.calls += 1 + t = float(timestep.reshape(-1)[0]) + self.call_timesteps.append(t) + # smooth damped-oscillatory velocity trajectory (the feature-ODE class) + return self.base * math.exp(-1.5 * t) * math.cos(3.0 * t) + 0.3 * t + + +class MockPipeline: + """Minimal stand-in for Hunyuan3DDiTFlowMatchingPipeline's denoise loop.""" + + num_train_timesteps = 1000 + + def __init__(self, model): + self.model = model + + def __call__(self, num_inference_steps=30, guidance_scale=5.5): + sigmas = torch.linspace(0, 1, num_inference_steps) + timesteps = sigmas * self.num_train_timesteps # increasing, starts at 0 + latents = torch.randn(1, 8, 4) + outputs = [] + for t in timesteps: + latent_model_input = torch.cat([latents] * 2) + timestep = t.expand(latent_model_input.shape[0]) / self.num_train_timesteps + noise_pred = self.model(latent_model_input, timestep, cond=None, guidance=None) + cond_pred, uncond_pred = noise_pred.chunk(2) + noise_pred = uncond_pred + guidance_scale * (cond_pred - uncond_pred) + latents = latents - 0.01 * noise_pred # toy scheduler step + outputs.append(noise_pred.clone()) + return latents, outputs + + +def expected_compute_steps(num_steps, interval, warmup): + """Reference schedule: hicache_pp computes warmup steps, then 1 of every + `interval` (counter-based: a compute step resets the skip counter). The + patch floors warmup at 1 — the first step of a run has nothing to forecast + from.""" + warmup = max(1, warmup) + computes, counter = [], 0 + for s in range(num_steps): + if s < warmup or counter >= interval - 1: + computes.append(s) + counter = 0 + else: + counter += 1 + return computes + + +# --------------------------------------------------------------------------- +# schedule / skip behaviour +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("method", ["hermite", "dmd", "auto"]) +@pytest.mark.parametrize("interval,warmup", [(3, 2), (5, 2), (4, 0)]) +def test_compute_calls_follow_schedule(method, interval, warmup): + torch.manual_seed(0) + steps = 30 + dit = MockDiT() + pipe = MockPipeline(dit) + hp.apply_hicache(pipe, method=method, interval=interval, warmup_steps=warmup) + + pipe(num_inference_steps=steps) + + expected = expected_compute_steps(steps, interval, warmup) + assert dit.calls == len(expected), ( + f"DiT ran {dit.calls} times, schedule says {len(expected)}") + assert pipe.model.computed_steps == len(expected) + assert pipe.model.skipped_steps == steps - len(expected) + # the real forwards happened exactly at the scheduled timesteps + sigmas = torch.linspace(0, 1, steps) + assert dit.call_timesteps == pytest.approx([float(sigmas[i]) for i in expected]) + + +class ConstantDiT(MockDiT): + """Constant velocity — every forecast basis is exact on this class + (all finite differences vanish), independent of sign/scale conventions.""" + + def forward(self, latent_model_input, timestep, cond, guidance=None): + self.calls += 1 + t = float(timestep.reshape(-1)[0]) + self.call_timesteps.append(t) + return self.base.clone() + + +class ExpDiT(MockDiT): + """Velocity exponential in the timestep: base * exp(-a*t) — the DMD/Prony + solution class, on which the DMD forecast is exact once its uniform + snapshot window (4) has filled.""" + + def forward(self, latent_model_input, timestep, cond, guidance=None): + self.calls += 1 + t = float(timestep.reshape(-1)[0]) + self.call_timesteps.append(t) + return self.base * math.exp(-2.0 * t) + + +def _run_pair(ref_dit, acc_dit, method, steps=30, interval=4, warmup=2, **kw): + acc_dit.load_state_dict(ref_dit.state_dict()) + 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) + torch.manual_seed(1) + _, acc_outputs = pipe(num_inference_steps=steps) + computes = expected_compute_steps(steps, interval, warmup) + skipped = [s for s in range(steps) if s not in computes] + assert skipped, "test must exercise skipped steps" + return ref_outputs, acc_outputs, computes, skipped + + +@pytest.mark.parametrize("method", ["hermite", "dmd", "auto"]) +def test_forecast_exact_on_constant_trajectory(method): + """Convention-free exactness check: on a constant velocity every skipped + step must reproduce the real DiT output bit-for-bit (all finite + differences vanish; the cached anchor is the exact forecast).""" + torch.manual_seed(0) + ref_outputs, acc_outputs, _, skipped = _run_pair( + ConstantDiT(), ConstantDiT(), method) + for s in skipped: + assert torch.allclose(acc_outputs[s], ref_outputs[s]), \ + f"step {s}: forecast not exact on constant series" + + +def test_patch_reproduces_hicache_pp_forecast_bitwise(): + """The node's claim is that it faithfully ports the hunyuan2-plus-plus + pipeline wiring of hicache-pp. Replay the patched run's compute anchors + through hicache_pp's own state machine and check that the patch served + the library's forecast bit-for-bit on every skipped step.""" + from hicache_pp import (hicache_init, hicache_decide, + hicache_update_derivatives, hicache_forecast) + + torch.manual_seed(0) + steps, interval, warmup = 30, 4, 2 + dit_a = MockDiT() + dit_b = MockDiT() + dit_b.load_state_dict(dit_a.state_dict()) + + pipe = MockPipeline(dit_a) + hp.apply_hicache(pipe, method="hermite", interval=interval, warmup_steps=warmup) + torch.manual_seed(1) + _, acc_outputs = pipe(num_inference_steps=steps) + + # reference: drive hicache_pp directly, exactly like the adapter's loop + state = hicache_init(num_steps=steps, interval=interval, max_order=1, + first_enhance=warmup, end_enhance=steps + 1, sigma=0.5) + sigmas = torch.linspace(0, 1, steps) + latents = torch.randn(1, 8, 4) # same shapes; model output ignores latents + for s in range(steps): + if hicache_decide(state) == "forecast": + ref_raw = hicache_forecast(state) + else: + lmi = torch.cat([latents] * 2) + timestep = sigmas[s].expand(lmi.shape[0]) + ref_raw = dit_b(lmi, timestep, cond=None) + hicache_update_derivatives(state, ref_raw.detach()) + state["step"] += 1 + # same CFG combine the pipeline loop applies after the model call + cond_pred, uncond_pred = ref_raw.chunk(2) + ref = uncond_pred + 5.5 * (cond_pred - uncond_pred) + assert torch.equal(acc_outputs[s], ref), \ + f"step {s}: patch output diverges from direct hicache_pp wiring" + + +def test_dmd_forecast_exact_on_exponential_trajectory(): + """DMD is exact on its solution class (sums of exponentials) once 4 + uniformly spaced snapshots exist; before that it falls back to Hermite + (which only needs to stay finite here).""" + torch.manual_seed(0) + steps, interval, warmup = 30, 4, 2 + ref_outputs, acc_outputs, computes, skipped = _run_pair( + ExpDiT(), ExpDiT(), "dmd", steps=steps, interval=interval, warmup=warmup) + # 4th uniformly spaced compute step (spacing=interval) closes the window + uniform = [c for c in computes if c >= warmup] + window_full_at = uniform[3] + checked = 0 + for s in skipped: + assert torch.isfinite(acc_outputs[s]).all() + if s > window_full_at: + rel = float((acc_outputs[s] - ref_outputs[s]).norm() + / ref_outputs[s].norm().clamp_min(1e-9)) + assert rel < 1e-3, f"step {s}: DMD not exact on exponential ({rel})" + checked += 1 + assert checked >= 6, "too few post-window skipped steps exercised" + + +def test_skipped_outputs_finite_and_better_than_naive_reuse(): + """On a smooth exponential trajectory the DMD forecast must be finite on + every skipped step and, in aggregate, beat the naive 'reuse the last + computed output' cache (what plain output caching would do).""" + torch.manual_seed(0) + steps, interval, warmup = 30, 4, 2 + ref_outputs, acc_outputs, computes, skipped = _run_pair( + ExpDiT(), ExpDiT(), "dmd", steps=steps, interval=interval, warmup=warmup) + + # every skipped step must be forecast-filled with a finite tensor + for s in skipped: + assert torch.isfinite(acc_outputs[s]).all() + + # once the DMD snapshot window has filled (4 uniformly spaced computes), + # the forecast must beat the naive 'reuse last computed output' cache + window_full_at = [c for c in computes if c >= warmup][3] + forecast_err, reuse_err = 0.0, 0.0 + for s in (s for s in skipped if s > window_full_at): + last_compute = max(c for c in computes if c < s) + forecast_err += float((acc_outputs[s] - ref_outputs[s]).norm()) + reuse_err += float((ref_outputs[last_compute] - ref_outputs[s]).norm()) + assert reuse_err > 0 + assert forecast_err < 0.1 * reuse_err, ( + f"forecast ({forecast_err:.4f}) must beat naive reuse ({reuse_err:.4f})") + + # compute steps are bit-identical to the reference model output + for s in computes: + assert torch.allclose(acc_outputs[s], ref_outputs[s]) + + +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(num_inference_steps=20) + assert dit.calls == 20 and pipe.model.skipped_steps == 0 + + +# --------------------------------------------------------------------------- +# method plumbing: the right forecaster actually runs +# --------------------------------------------------------------------------- +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(num_inference_steps=40) + st = pipe.model._state + assert st["backend"] == "dmd" + assert len(st["dmd_snapshots"]) >= 4, "DMD snapshot window never filled" + assert pipe.model.skipped_steps > 0 + + +def test_hermite_method_keeps_snapshots_empty(): + dit = MockDiT() + pipe = MockPipeline(dit) + hp.apply_hicache(pipe, method="hermite", interval=4, warmup_steps=2) + pipe(num_inference_steps=40) + st = pipe.model._state + assert st["backend"] == "hermite" + assert st["dmd_snapshots"] == [] + assert st["derivatives"], "Hermite anchors missing" + + +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(num_inference_steps=40) + st = pipe.model._state + assert st["backend"] == "auto" + # auto caches its per-compute-step holdout selection once enough snapshots exist + assert st.get("_auto_choice") in ("dmd", "hermite") + + +# --------------------------------------------------------------------------- +# run-boundary reset + passthrough +# --------------------------------------------------------------------------- +def test_state_resets_between_sampling_runs(): + dit = MockDiT() + pipe = MockPipeline(dit) + hp.apply_hicache(pipe, method="dmd", interval=4, warmup_steps=2) + + pipe(num_inference_steps=20) + first_run_calls = dit.calls + pipe(num_inference_steps=20) # timestep drops back to 0 -> reset + + expected = len(expected_compute_steps(20, 4, 2)) + assert first_run_calls == expected + assert dit.calls == 2 * expected, "second run must re-warm, not inherit state" + assert pipe.model.computed_steps == expected # stats are per-run + + +def test_attribute_and_device_passthrough(): + dit = MockDiT() + pipe = MockPipeline(dit) + 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) + assert pipe.model.inner is dit + + +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 + + +# --------------------------------------------------------------------------- +# parameter validation +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("kwargs,match", [ + (dict(method="taylor"), "method"), + (dict(interval=0), "interval"), + (dict(warmup_steps=-1), "warmup_steps"), + (dict(max_order=0), "max_order"), + (dict(sigma=1.0), "sigma"), + (dict(sigma=0.0), "sigma"), + (dict(dmd_history=2), "dmd_history"), +]) +def test_bad_params_rejected(kwargs, match): + base = dict(method="dmd", interval=4, warmup_steps=2, + max_order=1, sigma=0.5, dmd_history=5) + base.update(kwargs) + with pytest.raises(ValueError, match=match): + hp.validate_config(**base) + + +def test_apply_requires_model_attribute(): + with pytest.raises(TypeError, match="no `.model` attribute"): + hp.apply_hicache(object(), method="dmd")