fix: support guard-then-load lazy pipelines

.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

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

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():

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)