wip: release readiness pass

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.

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",

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,
         ),)
 
 

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]

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"])
+

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")