"""Standalone unit tests for the TRELLIS-HiCache patch logic. No ComfyUI, no GPU, no TRELLIS: a dummy DiT and a fake SparseTensor stand in for the real models, so the scheduling / CFG-routing / run-boundary / copy-on-patch logic is checked deterministically on CPU. """ import sys, os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import torch import pytest from trellis_hicache_patch import ( HiCacheModelPatch, apply_hicache, remove_hicache, validate_config, ) class DummyDiT(torch.nn.Module): """Counts forward calls; returns a deterministic dense velocity.""" def __init__(self, dim=8): super().__init__() self.dim = dim self.calls = 0 self.lin = torch.nn.Linear(dim, dim) def forward(self, x, t, cond=None, **kw): self.calls += 1 # smooth-in-t signal so the Hermite forecast is accurate tv = float(t.reshape(-1)[0]) if torch.is_tensor(t) else float(t) return torch.ones(1, self.dim) * tv class FakeSparse: """Minimal stand-in for a TRELLIS SparseTensor: .feats + .replace().""" def __init__(self, feats): self.feats = feats def replace(self, feats): return FakeSparse(feats) def _trellis_t_seq(steps=25): # TRELLIS scales t to [0, 1000], decreasing. import numpy as np return [float(x) for x in np.linspace(1, 0, steps + 1)[:-1] * 1000.0] def test_validate_config_rejects_bad_params(): with pytest.raises(ValueError): validate_config("nope", 3, 2, 1, 0.5, 5) with pytest.raises(ValueError): validate_config("hermite", 0, 2, 1, 0.5, 5) with pytest.raises(ValueError): validate_config("hermite", 3, 2, 1, 1.5, 5) with pytest.raises(ValueError, match="dmd_history"): validate_config("dmd", 3, 2, 1, 0.5, 3) with pytest.raises(ValueError, match="dmd_history"): validate_config("auto", 3, 2, 1, 0.5, 4) def test_skips_steps_on_decreasing_trellis_schedule(): dit = DummyDiT() patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2) for t in _trellis_t_seq(25): patch(torch.zeros(1, 8), t) # 25 steps, interval 3: clearly fewer real forwards than steps, some skipped. assert patch.skipped_steps > 0 assert dit.calls == patch.computed_steps assert patch.computed_steps + patch.skipped_steps == 25 def test_split_cfg_routes_two_states(): """Each timestep called twice (cond, uncond) must use two forecast states and still skip; the DiT must be called far fewer than 2x steps.""" dit = DummyDiT() patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2) for t in _trellis_t_seq(25): patch(torch.zeros(1, 8), t) # cond patch(torch.zeros(1, 8), t) # uncond (repeated t) assert patch.skipped_steps > 0 assert patch.computed_steps + patch.skipped_steps == 50 assert dit.calls < 50 # genuinely skipped DiT forwards def test_explicit_run_and_branch_identity_isolation_for_same_timestep_retry(): """An integration that knows job/CFG identity can disambiguate equal t.""" dit = DummyDiT() patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=1) t0 = torch.tensor([1000.0]) x = torch.zeros(1, 8) patch(x, t0, hicache_run_id="run-a", hicache_branch_id="cond") assert patch.run_id == "run-a" assert patch.branch_id == "cond" patch(x, t0, hicache_run_id="run-a", hicache_branch_id="uncond") assert patch.branch_id == "uncond" assert patch.telemetry["branches"]["cond"]["decisions"]["full"] == 1 assert patch.telemetry["branches"]["uncond"]["decisions"]["full"] == 1 # A new retry with the same one-step t=0-equivalent timestep must not inherit # either CFG branch's anchor or counter. patch(x, t0, hicache_run_id="run-b", hicache_branch_id="cond") assert patch.run_id == "run-b" assert patch.computed_steps == 1 and patch.skipped_steps == 0 assert dit.calls == 3 def test_telemetry_reports_fallbacks_and_is_detached(): dit = DummyDiT() patch = HiCacheModelPatch(dit, method="dmd", interval=3, warmup_steps=2, dmd_history=4, stage_id="slat_flow_model") for t in _trellis_t_seq(35): patch(torch.zeros(1, 8), t) telemetry = patch.telemetry assert telemetry["run_id"] == patch.run_id assert telemetry["stage_id"] == "slat_flow_model" assert telemetry["cfg_mode"] == "split_or_single" assert telemetry["method_counts"]["hermite"] > 0 assert telemetry["method_counts"]["dmd"] > 0 assert sum(telemetry["fallbacks"].values()) > 0 # Central hicache telemetry is returned detached from the live state. telemetry["decisions"]["full"] = -1 assert patch.telemetry["decisions"]["full"] >= 1 def test_new_run_resets_on_direction_reversal(): dit = DummyDiT() patch = HiCacheModelPatch(dit, method="hermite", interval=3, warmup_steps=2) for t in _trellis_t_seq(10): patch(torch.zeros(1, 8), t) first = patch.computed_steps + patch.skipped_steps # a second run restarts t near 1000 (direction flips up) -> reset for t in _trellis_t_seq(10): patch(torch.zeros(1, 8), t) assert patch.computed_steps + patch.skipped_steps == 10 # counters reset, not 20 assert first == 10 def test_sparse_output_forecast_rebuilds_sparse(): """SLaT-style sparse output: forecast must return a FakeSparse, not a tensor.""" class SparseDiT(torch.nn.Module): def __init__(self): super().__init__() self.calls = 0 def forward(self, x, t, cond=None, **kw): self.calls += 1 tv = float(t.reshape(-1)[0]) if torch.is_tensor(t) else float(t) return FakeSparse(torch.ones(4, 8) * tv) patch = HiCacheModelPatch(SparseDiT(), method="hermite", interval=2, warmup_steps=1) outs = [patch(None, t) for t in _trellis_t_seq(12)] assert all(isinstance(o, FakeSparse) for o in outs) assert patch.skipped_steps > 0 def test_apply_remove_is_copy_on_patch(): class Pipe: def __init__(self): self.models = {"sparse_structure_flow_model": DummyDiT(), "slat_flow_model": DummyDiT()} p = Pipe() patched = apply_hicache(p, method="hermite", interval=2, stages="both") # original untouched assert not getattr(p.models["slat_flow_model"], "_hicache_is_patch", False) assert getattr(patched.models["slat_flow_model"], "_hicache_is_patch", False) assert patched.models is not p.models # remove restores clean = remove_hicache(patched) assert not getattr(clean.models["slat_flow_model"], "_hicache_is_patch", False) def test_apply_never_nests(): class Pipe: def __init__(self): self.models = {"sparse_structure_flow_model": DummyDiT(), "slat_flow_model": DummyDiT()} p = Pipe() a = apply_hicache(p, interval=2, stages="both") b = apply_hicache(a, interval=3, stages="both") # re-patch inner = b.models["slat_flow_model"].inner assert not getattr(inner, "_hicache_is_patch", False) # unwrapped, not nested assert b.models["slat_flow_model"].interval == 3 def test_wrap_none_does_not_crash_on_attr_access(): """Regression: GGUF/lazy pipelines leave pipeline.models[key]=None at patch time. Wrapping None must not raise, and an unknown-attribute lookup must give a clear 'not loaded yet' error, NOT 'object has no attribute inner'.""" patch = HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2) assert patch.inner is None with pytest.raises(AttributeError) as ei: _ = patch.some_config_attr msg = str(ei.value) assert "not loaded yet" in msg and "some_config_attr" in msg with pytest.raises(RuntimeError): patch(torch.zeros(1, 8), _trellis_t_seq(4)[0]) def test_bind_inner_materializes_lazy_model(): patch = HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2) dit = DummyDiT() assert patch.bind_inner(dit) is patch assert patch.inner is dit assert "inner" in patch._modules for t in _trellis_t_seq(25): patch(torch.zeros(1, 8), t) assert patch.skipped_steps > 0 assert dit.calls == patch.computed_steps assert patch.computed_steps + patch.skipped_steps == 25 def test_lazy_pipeline_autobinds_on_assignment(): class LazyPipe: def __init__(self): self.models = {"sparse_structure_flow_model": None, "slat_flow_model": None} p = LazyPipe() patched = apply_hicache(p, method="hermite", interval=3, warmup_steps=2, stages="sparse_structure") slot = patched.models["sparse_structure_flow_model"] assert getattr(slot, "_hicache_is_patch", False) assert slot.inner is None dit = DummyDiT() patched.models["sparse_structure_flow_model"] = dit # sampler materializes it assert patched.models["sparse_structure_flow_model"] is slot assert slot.inner is dit for t in _trellis_t_seq(25): slot(torch.zeros(1, 8), t) assert slot.skipped_steps > 0 assert dit.calls == slot.computed_steps def test_eager_inner_still_registered_as_submodule(): dit = DummyDiT() patch = HiCacheModelPatch(dit, method="hermite", interval=3) assert patch.inner is dit assert "inner" in patch._modules assert any(k.startswith("inner.") for k in patch.state_dict().keys()) assert any(p is q for q in dit.parameters() for p in patch.parameters()) def test_stages_selector(): class Pipe: def __init__(self): self.models = {"sparse_structure_flow_model": DummyDiT(), "slat_flow_model": DummyDiT()} p = Pipe() ss = apply_hicache(p, stages="sparse_structure") assert getattr(ss.models["sparse_structure_flow_model"], "_hicache_is_patch", False) assert not getattr(ss.models["slat_flow_model"], "_hicache_is_patch", False) slat = apply_hicache(p, stages="slat") assert not getattr(slat.models["sparse_structure_flow_model"], "_hicache_is_patch", False) assert slat.models["slat_flow_model"].stage_id == "slat_flow_model" both = apply_hicache(p, stages="both") assert both.models["sparse_structure_flow_model"].stage_id == "sparse_structure_flow_model" assert both.models["slat_flow_model"].stage_id == "slat_flow_model"