From fd2c6a5fdaeeabce3bb508bb7bf9f446133271ea Mon Sep 17 00:00:00 2001 From: Krishi Attri Date: Sat, 25 Jul 2026 00:34:10 +0000 Subject: [PATCH] fix: support guard-then-load lazy pipelines --- .github/workflows/ci.yml | 28 ++++++++++++++++++ pyproject.toml | 6 ++++ tests/test_patch.py | 58 +++++++++++++++++++++++++++--------- trellis_hicache_patch.py | 63 ++++++++++++++++++++++++++++++---------- 4 files changed, 127 insertions(+), 28 deletions(-) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..d747c7a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,28 @@ +name: CI + +on: + push: + branches: [master] + pull_request: + branches: [master] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install package and test dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[test]" + - name: Run CPU unit tests + run: python -m pytest -q + - name: Build wheel from clean metadata + run: python -m pip wheel . --no-deps --wheel-dir dist diff --git a/pyproject.toml b/pyproject.toml index 7958d14..6d13343 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,9 @@ requires-python = ">=3.9" dependencies = ["hicache-pp>=1.1.0"] classifiers = ["Development Status :: 4 - Beta"] +[project.optional-dependencies] +test = ["pytest>=8", "numpy>=1.24", "torch>=2.1"] + [project.urls] Repository = "https://github.com/Archerkattri/ComfyUI-TRELLIS2-HiCache" BugTracker = "https://github.com/Archerkattri/ComfyUI-TRELLIS2-HiCache/issues" @@ -23,3 +26,6 @@ Icon = "https://raw.githubusercontent.com/Archerkattri/ComfyUI-TRELLIS2-HiCache/ [tool.pytest.ini_options] addopts = "--import-mode=importlib" testpaths = ["tests"] + +[tool.setuptools] +py-modules = ["trellis_hicache_patch", "nodes"] diff --git a/tests/test_patch.py b/tests/test_patch.py index bace9ea..7490eec 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -162,30 +162,62 @@ def test_bind_inner_materializes_lazy_model(): assert patch.computed_steps + patch.skipped_steps == 25 -def test_lazy_pipeline_autobinds_on_assignment(): - """End-to-end lazy flow: patch a pipeline whose flow model is still None, then - the sampler assigns the real model into pipeline.models[key]; the patch must - bind it (not be overwritten) and then skip DiT steps.""" +def test_lazy_pipeline_guard_then_load_still_triggers(): + """A lazy loader's ``is None`` guard must still run, including after unload.""" class LazyPipe: def __init__(self): self.models = {"sparse_structure_flow_model": None, "shape_slat_flow_model_512": None, "shape_slat_flow_model_1024": None} + self.loads = [] + + def load_sparse_structure_model(self): + if self.models["sparse_structure_flow_model"] is None: + model = DummyDiT() + self.loads.append(model) + self.models["sparse_structure_flow_model"] = model + return self.models["sparse_structure_flow_model"] + p = LazyPipe() patched = apply_hicache(p, method="hermite", interval=3, warmup_steps=2, stages="sparse_structure") - slot = patched.models["sparse_structure_flow_model"] + assert patched.models["sparse_structure_flow_model"] is None + + slot = patched.load_sparse_structure_model() assert getattr(slot, "_hicache_is_patch", False) - assert slot.inner is None # deferred - # sampler materializes the real DiT (the exact GGUF lazy-load path) - dit = DummyDiT() - patched.models["sparse_structure_flow_model"] = dit - # same patch object, now bound -- not overwritten by the raw DiT - assert patched.models["sparse_structure_flow_model"] is slot - assert slot.inner is dit + assert slot.inner is patched.loads[0] for t in _trellis_t_seq(25): slot(torch.zeros(1, 8), t) assert slot.skipped_steps > 0 - assert dit.calls == slot.computed_steps + assert patched.loads[0].calls == slot.computed_steps + + patched.models["sparse_structure_flow_model"] = None + assert patched.models["sparse_structure_flow_model"] is None + replacement = patched.load_sparse_structure_model() + assert len(patched.loads) == 2 + assert replacement is not slot + assert getattr(replacement, "_hicache_is_patch", False) + assert replacement.inner is patched.loads[1] + + +def test_remove_hicache_clears_pending_before_any_load(): + """Disabling a pending patch must not re-enable it on a later assignment.""" + class LazyPipe: + def __init__(self): + self.models = {"sparse_structure_flow_model": None, + "shape_slat_flow_model_512": None, + "shape_slat_flow_model_1024": None} + + patched = apply_hicache( + LazyPipe(), method="hermite", interval=3, warmup_steps=2, + stages="sparse_structure", + ) + clean = remove_hicache(patched) + assert type(clean.models) is dict + assert clean.models["sparse_structure_flow_model"] is None + + dit = DummyDiT() + clean.models["sparse_structure_flow_model"] = dit + assert clean.models["sparse_structure_flow_model"] is dit def test_eager_inner_still_registered_as_submodule(): diff --git a/trellis_hicache_patch.py b/trellis_hicache_patch.py index 68508cc..63521e5 100644 --- a/trellis_hicache_patch.py +++ b/trellis_hicache_patch.py @@ -315,18 +315,38 @@ class HiCacheModelPatch(torch.nn.Module): class _LazyPatchDict(dict): - """A ``pipeline.models`` drop-in that keeps patches alive across lazy loads. + """A ``pipeline.models`` drop-in that patches models after lazy loads. GGUF / lazy Trellis2 pipelines leave a flow-model slot empty (``None``) at patch time and assign the real DiT later, from inside the sampler, via - ``pipeline.models[key] = model``. With a plain dict that assignment would - overwrite the :class:`HiCacheModelPatch` wrapper and silently disable the - cache. This subclass instead *binds* the freshly loaded model into the - existing patch, so acceleration survives the lazy load with no user action. - It is a ``dict`` subclass, so ``isinstance(models, dict)`` checks still pass. + ``pipeline.models[key] = model``. The slot must remain the real ``None`` + singleton until then because loaders commonly guard the assignment with + ``if pipeline.models[key] is None``. Pending patch configuration therefore + lives out-of-band in ``_pending``; assigning a real model wraps it at that + boundary. Pending entries survive assignment back to ``None`` so VRAM + unload/reload cycles receive a fresh patch each time. """ + def __init__(self, *args, **kwargs): + source = args[0] if args else None + super().__init__(*args, **kwargs) + inherited = getattr(source, "_pending", {}) + self._pending = {key: dict(config) for key, config in inherited.items()} + + def set_pending(self, key: str, config: Dict[str, Any]) -> None: + self._pending[key] = dict(config) + dict.__setitem__(self, key, None) + def __setitem__(self, key, value): + if key in self._pending: + if value is None: + dict.__setitem__(self, key, None) + return + if not getattr(value, "_hicache_is_patch", False): + value = HiCacheModelPatch(value, **self._pending[key]) + dict.__setitem__(self, key, value) + return + existing = self.get(key) if (isinstance(existing, HiCacheModelPatch) and value is not None @@ -369,24 +389,33 @@ def apply_hicache(pipeline: Any, *, method: str = "hermite", interval: int = 3, """ keys = _resolve_keys(pipeline, stages) patched = copy.copy(pipeline) - # _LazyPatchDict (a dict subclass) so a later pipeline.models[key] = model - # from a lazy/GGUF sampler binds into the patch instead of overwriting it. + # A dict subclass keeps real None sentinels visible to lazy/GGUF loader + # guards, then wraps the model at the assignment boundary. patched.models = _LazyPatchDict(pipeline.models) # copy so the original is untouched lazy_keys = [] for key in keys: inner = patched.models[key] if getattr(inner, "_hicache_is_patch", False): inner = inner.inner # replace, never nest + config = { + "method": method, + "interval": interval, + "warmup_steps": warmup_steps, + "max_order": max_order, + "sigma": sigma, + "dmd_history": dmd_history, + } if inner is None: lazy_keys.append(key) - dict.__setitem__(patched.models, key, HiCacheModelPatch( - inner, method=method, interval=interval, warmup_steps=warmup_steps, - max_order=max_order, sigma=sigma, dmd_history=dmd_history, - )) + patched.models.set_pending(key, config) + else: + if key in patched.models._pending: + patched.models._pending[key] = dict(config) + dict.__setitem__(patched.models, key, HiCacheModelPatch(inner, **config)) if lazy_keys: logger.warning( - "[TRELLIS-HiCache] %s not loaded yet (lazy/GGUF pipeline); wrapped a " - "deferred patch on %s -- it binds when the sampler materializes the model.", + "[TRELLIS-HiCache] %s not loaded yet (lazy/GGUF pipeline); queued a " + "deferred patch on %s -- assignment wraps the materialized model.", type(pipeline).__name__, lazy_keys, ) logger.info( @@ -400,7 +429,11 @@ def remove_hicache(pipeline: Any) -> Any: """Return ``pipeline`` with the original DiTs restored (copy-on-unpatch).""" if not hasattr(pipeline, "models") or not isinstance(pipeline.models, dict): return pipeline - if not any(getattr(m, "_hicache_is_patch", False) for m in pipeline.models.values()): + has_patch = any( + getattr(m, "_hicache_is_patch", False) for m in pipeline.models.values() + ) + has_pending = bool(getattr(pipeline.models, "_pending", {})) + if not has_patch and not has_pending: return pipeline clean = copy.copy(pipeline) clean.models = dict(pipeline.models)