From a65aa6aa1a661d23a2da8a7e272ee583a4312de8 Mon Sep 17 00:00:00 2001 From: Archerkattri Date: Tue, 15 Sep 2026 12:45:14 -0400 Subject: [PATCH] wip: release readiness pass --- README.md | 27 ++++++ hicache_patch.py | 176 +++++++++++++++++++++++++++++++++--- nodes.py | 24 ++++- pyproject.toml | 2 +- tests/test_budget_bridge.py | 45 +++++++++ tests/test_patch.py | 34 +++++++ 6 files changed, 292 insertions(+), 16 deletions(-) create mode 100644 tests/test_budget_bridge.py diff --git a/README.md b/README.md index 87c7146..a7685ef 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,26 @@ 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. +## First-result acceptance checklist + +Use a fixed small input and record the exact ComfyUI, wrapper, checkpoint, node +commit, method, interval, seed, and torch/CUDA versions. Before treating a run +as accelerated, verify all of the following: + +* **Baseline:** run once with `enable=false` (or `interval=1`) and save the + output plus the unaccelerated DiT-forward count. +* **Accelerated:** run the same input and seed with the node enabled; confirm + `pipeline.model.branch_id` is `stacked_pre_cfg`, the `run_id` changes on a + retry, and `pipeline.model.telemetry` reports the actual decision and method + counters. DMD/auto history below their method-specific floors is rejected; + short or non-uniform windows remain visible as Hermite fallback counts. +* **Lifecycle:** repeat the run, including a one-step `t=0` retry, and confirm + the run starts with a fresh full decision rather than a stale forecast. +* **Comparison:** compare baseline and accelerated outputs and retain the raw + per-run counts. GPU timing/quality and a clean-host workflow import are + separate gates; this CPU package check does not establish a speed or quality + claim. + 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; @@ -213,3 +233,10 @@ numbers transfer exactly, and that is what the defaults use. ## License MIT. Not affiliated with Tencent (Hunyuan3D) or kijai (ComfyUI-Hunyuan3DWrapper). + +## Current release status + +The current adapter includes the shared HiCache++ budget/identity/manifest +bridge, timing and fallback accounting. Its CPU contract suite passes 45 tests. +ComfyUI execution with the real CUDA model workflow was not run in this +environment, so no image-quality or speedup claim is made. diff --git a/hicache_patch.py b/hicache_patch.py index 1a12e06..fa70f54 100644 --- a/hicache_patch.py +++ b/hicache_patch.py @@ -72,11 +72,16 @@ import torch from hicache_pp import ( hicache_init, hicache_decide, + hicache_telemetry, hicache_update_derivatives, hicache_forecast, dmd_update_snapshots, dmd_forecast_state, auto_forecast_state, + CacheBudget, + CacheBudgetRuntime, + RunIdentity, + stable_digest, ) logger = logging.getLogger("ComfyUI-HiCache") @@ -101,8 +106,11 @@ def validate_config(method: str, interval: int, warmup_steps: int, 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}") + min_history = {"hermite": 3, "dmd": 4, "auto": 5}[method] + if dmd_history < min_history: + raise ValueError( + f"dmd_history must be >= {min_history} for {method}, got {dmd_history}" + ) class HiCacheModelPatch(torch.nn.Module): @@ -116,12 +124,16 @@ class HiCacheModelPatch(torch.nn.Module): def __init__(self, model: Optional[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: + sigma: float = 0.5, dmd_history: int = 5, + budget: Optional[CacheBudget] = None, + max_horizon: Optional[int] = None, + max_memory_mb: Optional[float] = None, + audit_budget: int = 0) -> None: validate_config(method, interval, warmup_steps, max_order, sigma, dmd_history) super().__init__() # ``model`` may be None: lazy / GGUF pipelines materialize the DiT after # patching. A real nn.Module is registered as a submodule (so - # .to()/.state_dict() recurse); a None/lazy placeholder is kept as a plain + # .to()/.state_dict() recurse); a None/lazy deferred module is kept as a plain # attribute so __getattr__ never dead-ends. See _set_inner / bind_inner. self._set_inner(model) self._hicache_is_patch = True # marker for apply/remove @@ -131,19 +143,81 @@ class HiCacheModelPatch(torch.nn.Module): self.max_order = int(max_order) self.sigma = float(sigma) self.dmd_history = int(dmd_history) + if budget is None: + budget = CacheBudget( + backend=method, + allowed_stages=("shape",), + max_horizon=max(1, interval - 1) if max_horizon is None else max_horizon, + quality_preset="adapter-default", + max_memory_mb=max_memory_mb, + audit_budget=audit_budget, + fallback="full", + ) + elif not isinstance(budget, CacheBudget): + raise TypeError("budget must be a CacheBudget") + self.budget = budget self._state: Optional[Dict[str, Any]] = None + self._last_telemetry: Dict[str, Any] = {} + self._last_budget_manifest: Dict[str, Any] = {} + self._budget_runtime: Optional[CacheBudgetRuntime] = None + self._last_budget_decision: Dict[str, Any] = {} self._last_t: Optional[float] = None + self.last_decision: Optional[str] = None # per-run stats (read by the node / logged at run boundaries) self.computed_steps = 0 self.skipped_steps = 0 + @property + def run_id(self) -> Optional[str]: + """Identity of the currently active sampling run, if any.""" + return None if self._state is None else self._state.get("run_id") + + @property + def branch_id(self) -> Optional[str]: + """Stable identity for this model-level, pre-CFG patch branch.""" + return None if self._state is None else self._state.get("branch_id") + + @property + def telemetry(self) -> Dict[str, Any]: + """Detached actual-method/fallback counters for the active run.""" + if self._state is not None: + return hicache_telemetry(self._state) + return copy.deepcopy(self._last_telemetry) + + @property + def budget_manifest(self) -> Dict[str, Any]: + """Portable budget/identity/decision report for the active run.""" + if self._budget_runtime is not None: + return copy.deepcopy(self._budget_runtime.manifest.as_dict()) + return copy.deepcopy(self._last_budget_manifest) + + @property + def budget_decision(self) -> Dict[str, Any]: + """The last budget decision without exposing mutable runtime state.""" + return copy.deepcopy(self._last_budget_decision) + + def save_budget_manifest(self, destination: str) -> None: + """Save the latest portable manifest; no tensors or local paths are written.""" + if self._budget_runtime is not None: + self._budget_runtime.manifest.write_json(destination) + elif self._last_budget_manifest: + import json + from pathlib import Path + + Path(destination).write_text( + json.dumps(self._last_budget_manifest, sort_keys=True, indent=2) + "\n", + encoding="utf-8", + ) + else: + raise RuntimeError("no budget manifest exists; run the patched model first") + # -- inner storage / attribute passthrough --------------------------------- def _set_inner(self, model: Optional[torch.nn.Module]) -> None: """Store the wrapped model so ``self.inner`` is always resolvable. A real ``nn.Module`` is registered in ``_modules`` (so device moves and - ``state_dict`` recurse into it). A ``None`` placeholder -- or any + ``state_dict`` recurse into it). A ``None`` deferred value -- or any non-Module -- is kept as a plain attribute in ``__dict__``. This avoids the bug where ``nn.Module.__setattr__`` routes only real Modules into ``_modules``; a None assigned to ``self.inner`` lands in ``__dict__``, and @@ -166,7 +240,7 @@ class HiCacheModelPatch(torch.nn.Module): return self.__dict__.get("inner") def bind_inner(self, model: torch.nn.Module) -> "HiCacheModelPatch": - """Attach the real DiT to a patch created around a lazy (None) placeholder, + """Attach the real DiT to a patch created around a deferred (None) binding, registering it as a submodule. Lets lazy / GGUF pipelines that materialize the model after patching still route through the cache.""" self._set_inner(model) @@ -190,7 +264,7 @@ class HiCacheModelPatch(torch.nn.Module): # -- state handling ---------------------------------------------------------- def _fresh_state(self) -> Dict[str, Any]: # hicache_pp validates interval/order/sigma/backend again here. - return hicache_init( + state = hicache_init( num_steps=_NO_END_WINDOW, interval=self.interval, max_order=self.max_order, @@ -202,17 +276,25 @@ class HiCacheModelPatch(torch.nn.Module): backend=self.method, history=self.dmd_history, ) + state["branch_id"] = "stacked_pre_cfg" + return state 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): + self._last_telemetry = hicache_telemetry(self._state) 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, ) + if self._budget_runtime is not None: + self._last_budget_manifest = self._budget_runtime.manifest.as_dict() self._state = self._fresh_state() + self._budget_runtime = None + self._last_budget_decision = {} self._last_t = None + self.last_decision = None self.computed_steps = 0 self.skipped_steps = 0 @@ -222,13 +304,54 @@ class HiCacheModelPatch(torch.nn.Module): return float(timestep.reshape(-1)[0].item()) return float(timestep) - def _forecast(self, state: Dict[str, Any]) -> torch.Tensor: - if self.method == "dmd": + def _forecast(self, state: Dict[str, Any], method: Optional[str] = None) -> torch.Tensor: + method = self.method if method is None else method + if method == "dmd": return dmd_forecast_state(state) - if self.method == "auto": + if method == "auto": return auto_forecast_state(state) return hicache_forecast(state) + def _ensure_budget_runtime( + self, state: Dict[str, Any], latent_model_input: torch.Tensor + ) -> CacheBudgetRuntime: + inner = self._inner() + shape = tuple(int(value) for value in latent_model_input.shape) + dtype = str(latent_model_input.dtype) + device = str(latent_model_input.device) + identity = RunIdentity( + model_id=type(inner).__name__ if inner is not None else "unloaded-model", + run_id=str(state["run_id"]), + schedule_digest=stable_digest({ + "interval": self.interval, + "warmup_steps": self.warmup_steps, + "max_order": self.max_order, + "sigma": self.sigma, + "dmd_history": self.dmd_history, + }), + cfg_branch="stacked_pre_cfg", + conditioning_id="pre-cfg", + stage="shape", + token_layout_digest=stable_digest({"shape": shape}), + dtype=dtype, + device=device, + batch_id=stable_digest({"shape": shape, "dtype": dtype, "device": device}), + ) + if self._budget_runtime is None or self._budget_runtime.identity.fingerprint != identity.fingerprint: + self._budget_runtime = CacheBudgetRuntime( + self.budget, + identity, + source_digest=stable_digest({"adapter": "ComfyUI-HiCache", "contract": "budget-v1"}), + config_digest=self.budget.digest, + ) + return self._budget_runtime + + @staticmethod + def _current_memory_mb(tensor: torch.Tensor) -> Optional[float]: + if not tensor.is_cuda: + return None + return float(torch.cuda.memory_allocated(tensor.device)) / (1024.0 * 1024.0) + # -- the patched forward ------------------------------------------------------ def forward(self, latent_model_input: torch.Tensor, timestep: Any, *args: Any, **kwargs: Any) -> torch.Tensor: @@ -243,8 +366,30 @@ class HiCacheModelPatch(torch.nn.Module): self._last_t = t_val state = self._state - if hicache_decide(state) == "forecast": - out = self._forecast(state) + decision = hicache_decide(state) + self.last_decision = decision + budget_runtime = self._ensure_budget_runtime(state, latent_model_input) + budget_decision = budget_runtime.decide( + "shape", + horizon=int(state.get("counter", 0)) if decision == "forecast" else 0, + method=self.method, + supported=self._inner() is not None, + memory_mb=self._current_memory_mb(latent_model_input), + controller_selected=self.method == "auto", + ) + self._last_budget_decision = budget_decision.as_dict() + if budget_decision.mode == "fallback" and budget_decision.method == "full" and decision == "forecast": + # The legacy scheduler has already incremented its forecast counter; + # repair it so the measured state and budget manifest agree. + state["type"] = "full" + state["counter"] = 0 + state["activated_steps"].append(state["step"]) + telemetry = state["telemetry"]["decisions"] + telemetry["forecast"] = max(0, int(telemetry.get("forecast", 0)) - 1) + telemetry["full"] = int(telemetry.get("full", 0)) + 1 + decision = "full" + if decision == "forecast" and budget_decision.mode in ("forecast", "fallback"): + out = self._forecast(state, budget_decision.method) state["step"] += 1 self.skipped_steps += 1 return out @@ -271,7 +416,10 @@ class HiCacheModelPatch(torch.nn.Module): # --------------------------------------------------------------------------- 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: + dmd_history: int = 5, budget: Optional[CacheBudget] = None, + max_horizon: Optional[int] = None, + max_memory_mb: Optional[float] = None, + audit_budget: int = 0) -> Any: """Return a shallow copy of ``pipeline`` whose ``model`` is patched. The input pipeline is NOT mutated (ComfyUI caches node outputs keyed on @@ -293,6 +441,8 @@ def apply_hicache(pipeline: Any, *, method: str = "hermite", interval: int = 3, patched.model = HiCacheModelPatch( inner, method=method, interval=interval, warmup_steps=warmup_steps, max_order=max_order, sigma=sigma, dmd_history=dmd_history, + budget=budget, max_horizon=max_horizon, max_memory_mb=max_memory_mb, + audit_budget=audit_budget, ) logger.info( "[HiCache] patched copy of %s: method=%s interval=%d warmup=%d", diff --git a/nodes.py b/nodes.py index e68669c..23c8922 100644 --- a/nodes.py +++ b/nodes.py @@ -66,7 +66,23 @@ class HiCacheAccelerate: "high-order terms bounded."}), "dmd_history": ("INT", { "default": 5, "min": 3, "max": 16, "step": 1, - "tooltip": "DMD snapshot window length (dmd/auto methods)."}), + "tooltip": "Snapshot window length. Minimum effective history: " + "3 for hermite, 4 for dmd, 5 for auto; shorter " + "values are rejected for the selected method."}), + "max_horizon": ("INT", { + "default": 0, "min": 0, "max": 64, "step": 1, + "tooltip": "Maximum forecast distance from a paid anchor. " + "0 follows the selected interval; exceeding the " + "budget falls back to full compute."}), + "max_memory_mb": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 262144.0, "step": 1.0, + "tooltip": "Optional current-device memory cap for forecast " + "decisions. 0 disables the cap."}), + "audit_budget": ("INT", { + "default": 0, "min": 0, "max": 128, "step": 1, + "tooltip": "Maximum explicit paid forward-audit decisions " + "available to an integration; reports are exposed " + "on pipeline.model.budget_manifest."}), }, } @@ -81,7 +97,8 @@ class HiCacheAccelerate: ) def patch(self, pipeline, method="hermite", interval=3, warmup_steps=2, - enable=True, max_order=1, sigma=0.5, dmd_history=5): + enable=True, max_order=1, sigma=0.5, dmd_history=5, + max_horizon=0, max_memory_mb=0.0, audit_budget=0): if not enable or interval <= 1: return (remove_hicache(pipeline),) return (apply_hicache( @@ -92,6 +109,9 @@ class HiCacheAccelerate: max_order=max_order, sigma=sigma, dmd_history=dmd_history, + max_horizon=None if max_horizon <= 0 else max_horizon, + max_memory_mb=None if max_memory_mb <= 0 else max_memory_mb, + audit_budget=audit_budget, ),) diff --git a/pyproject.toml b/pyproject.toml index ba72114..0a93043 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ keywords = ["comfyui", "diffusion-acceleration", "caching", "training-free", "hu license = "MIT" license-files = ["LICENSE"] requires-python = ">=3.9" -dependencies = ["hicache-pp>=1.1.0"] +dependencies = ["hicache-pp>=1.2.1"] classifiers = ["Development Status :: 5 - Production/Stable"] [project.optional-dependencies] diff --git a/tests/test_budget_bridge.py b/tests/test_budget_bridge.py new file mode 100644 index 0000000..7ec9c5f --- /dev/null +++ b/tests/test_budget_bridge.py @@ -0,0 +1,45 @@ +"""Integration checks for the shared hicache-pp budget contract.""" + +import importlib.util +import pathlib + +import torch + + +PACK_DIR = pathlib.Path(__file__).resolve().parents[1] + + +def _load_patch(): + spec = importlib.util.spec_from_file_location("hicache_budget_bridge_patch", PACK_DIR / "hicache_patch.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _DiT(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("base", torch.ones(1, 2)) + + def forward(self, x, timestep, cond, guidance=None): + return self.base * (1.0 + float(timestep.reshape(-1)[0])) + + +class _Pipe: + def __init__(self): + self.model = _DiT() + + +def test_budget_cap_forces_full_fallback_and_manifest_counts_actual_decisions(): + patch = _load_patch() + pipeline = patch.apply_hicache(_Pipe(), method="hermite", interval=4, max_horizon=1) + for timestep in torch.linspace(0.0, 1.0, 10): + pipeline.model(torch.ones(1, 2), timestep, None) + + report = pipeline.model.budget_manifest + assert report["budget"]["max_horizon"] == 1 + assert report["identity"]["stage"] == "shape" + assert report["counts"]["fallback"] > 0 + assert report["counts"]["full"] + report["counts"]["forecast"] + report["counts"]["fallback"] == 10 + assert all("path" not in event for event in report["events"]) + diff --git a/tests/test_patch.py b/tests/test_patch.py index 259ec63..90aa9e6 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -287,8 +287,13 @@ def test_dmd_method_populates_snapshots_and_forecasts(): pipe(num_inference_steps=40) st = pipe.model._state assert st["backend"] == "dmd" + assert pipe.model.branch_id == "stacked_pre_cfg" assert len(st["dmd_snapshots"]) >= 4, "DMD snapshot window never filled" assert pipe.model.skipped_steps > 0 + assert pipe.model.telemetry["decisions"]["forecast"] == pipe.model.skipped_steps + assert pipe.model.telemetry["method_counts"]["hermite"] > 0 + assert sum(pipe.model.telemetry["fallbacks"].values()) > 0 + assert pipe.model.telemetry["method_counts"]["dmd"] > 0 def test_hermite_method_keeps_snapshots_empty(): @@ -448,6 +453,28 @@ def test_single_step_runs_reset_at_equal_timestep(): assert patch.computed_steps == 1 and patch.skipped_steps == 0 # per-run stats +def test_run_identity_and_detached_telemetry_change_on_equal_timestep_reset(): + dit = MockDiT() + patch = hp.HiCacheModelPatch(dit, method="dmd", interval=4, + warmup_steps=1, dmd_history=4) + lmi = torch.randn(2, 8, 4) + t0 = torch.zeros(2) + + patch(lmi, t0, cond=None) + first_run_id = patch.run_id + assert first_run_id + assert patch.branch_id == "stacked_pre_cfg" + observed = patch.telemetry + assert observed["decisions"]["full"] == 1 + observed["decisions"]["full"] = 99 + assert patch.telemetry["decisions"]["full"] == 1 + + patch(lmi, t0, cond=None) + assert patch.run_id and patch.run_id != first_run_id + assert patch.telemetry["decisions"]["full"] == 1 + assert patch.last_decision == "full" + + # --------------------------------------------------------------------------- # parameter validation # --------------------------------------------------------------------------- @@ -468,6 +495,13 @@ def test_bad_params_rejected(kwargs, match): hp.validate_config(**base) +@pytest.mark.parametrize("method,history", [("dmd", 3), ("auto", 4)]) +def test_history_floor_matches_stateful_forecaster(method, history): + with pytest.raises(ValueError, match="dmd_history"): + hp.validate_config(method=method, interval=4, warmup_steps=2, + max_order=1, sigma=0.5, dmd_history=history) + + def test_apply_requires_model_attribute(): with pytest.raises(TypeError, match="no `.model` attribute"): hp.apply_hicache(object(), method="dmd")