test: lock lazy re-patch lifecycle

.github/workflows/ci.yml

@@ -24,7 +24,7 @@ jobs:
           pip install -e ".[test]"
       - name: Run CPU unit tests
         run: python -m pytest -q
-      - name: Run registry security style checks
-        run: python -m ruff check --select E702 .
+      - name: Run static checks
+        run: python -m ruff check .
       - name: Build wheel from clean metadata
         run: python -m pip wheel . --no-deps --wheel-dir dist

pyproject.toml

@@ -1,12 +1,17 @@
 # Packaging metadata for the registry.comfy.org Comfy Registry.
 # Field layout follows https://docs.comfy.org/registry/specifications
 
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
 [project]
 name = "comfyui-trellis2-hicache"
 description = "Training-free TRELLIS.2 image-to-3D acceleration: forecast the flow-matching velocity on skipped DiT steps (HiCache Hermite / HiCache++ DMD, via hicache-pp) across the sparse-structure and shape-SLaT stages (texture optional). ~2x faster, near-lossless. Pairs with visualbruno/ComfyUI-Trellis2."
 version = "0.1.3"
 keywords = ["comfyui", "diffusion-acceleration", "caching", "training-free", "trellis2", "image-to-3d"]
-license = { file = "LICENSE" }
+license = "MIT"
+license-files = ["LICENSE"]
 requires-python = ">=3.9"
 dependencies = ["hicache-pp>=1.1.0"]
 classifiers = ["Development Status :: 4 - Beta"]

tests/test_patch.py

@@ -4,7 +4,8 @@ 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
+import os
+import sys
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 import torch
@@ -201,6 +202,78 @@ def test_lazy_pipeline_guard_then_load_still_triggers():
     assert replacement.inner is patched.loads[1]
 
 
+def test_lazy_reapply_updates_config_across_unload_reload():
+    """Re-patching a loaded lazy slot must update its future reload config."""
+    class LazyPipe:
+        def __init__(self):
+            self.models = {
+                "sparse_structure_flow_model": None,
+                "shape_slat_flow_model_512": None,
+                "shape_slat_flow_model_1024": None,
+            }
+
+    first = apply_hicache(
+        LazyPipe(), method="hermite", interval=3, stages="sparse_structure",
+    )
+    first.models["sparse_structure_flow_model"] = DummyDiT()
+    assert first.models["sparse_structure_flow_model"].interval == 3
+
+    second = apply_hicache(
+        first, method="dmd", interval=5, stages="sparse_structure",
+    )
+    loaded = second.models["sparse_structure_flow_model"]
+    assert isinstance(loaded, HiCacheModelPatch)
+    assert loaded.interval == 5
+    assert loaded.method == "dmd"
+    assert second.models._pending["sparse_structure_flow_model"]["interval"] == 5
+    assert second.models._pending["sparse_structure_flow_model"]["method"] == "dmd"
+
+    second.models["sparse_structure_flow_model"] = None
+    second.models["sparse_structure_flow_model"] = DummyDiT()
+    reloaded = second.models["sparse_structure_flow_model"]
+    assert isinstance(reloaded, HiCacheModelPatch)
+    assert reloaded.interval == 5
+    assert reloaded.method == "dmd"
+
+
+def test_lazy_subset_reapply_preserves_unselected_pending_stages():
+    """A subset re-patch must not disable deferred wrapping for other stages."""
+    class LazyPipe:
+        def __init__(self):
+            self.models = {
+                "sparse_structure_flow_model": None,
+                "shape_slat_flow_model_512": None,
+                "shape_slat_flow_model_1024": None,
+                "tex_slat_flow_model_512": None,
+                "tex_slat_flow_model_1024": None,
+            }
+
+    all_stages = apply_hicache(
+        LazyPipe(), method="hermite", interval=3, stages="all",
+    )
+    all_stages.models["sparse_structure_flow_model"] = DummyDiT()
+
+    sparse_only = apply_hicache(
+        all_stages, method="dmd", interval=5, stages="sparse_structure",
+    )
+    assert set(sparse_only.models._pending) == set(all_stages.models)
+    assert sparse_only.models._pending["sparse_structure_flow_model"]["interval"] == 5
+    assert sparse_only.models._pending["shape_slat_flow_model_512"]["interval"] == 3
+    assert sparse_only.models._pending["tex_slat_flow_model_1024"]["method"] == "hermite"
+
+    sparse_only.models["shape_slat_flow_model_512"] = DummyDiT()
+    shape = sparse_only.models["shape_slat_flow_model_512"]
+    assert isinstance(shape, HiCacheModelPatch)
+    assert shape.interval == 3
+    assert shape.method == "hermite"
+
+    sparse_only.models["tex_slat_flow_model_1024"] = DummyDiT()
+    texture = sparse_only.models["tex_slat_flow_model_1024"]
+    assert isinstance(texture, HiCacheModelPatch)
+    assert texture.interval == 3
+    assert texture.method == "hermite"
+
+
 def test_remove_hicache_clears_pending_before_any_load():
     """Disabling a pending patch must not re-enable it on a later assignment."""
     class LazyPipe:

tests/validate_gpu.py

@@ -10,7 +10,10 @@ Run from a hermit-trellis2 checkout (which vendors the trellis2 package):
     PYTHONPATH=.:../comfyui-trellis2-hicache python ../comfyui-trellis2-hicache/tests/validate_gpu.py \
         --ckpt /home/krishi/workspace/data/weights/TRELLIS.2-4B --image <some.png>
 """
-import argparse, time, sys, types
+import argparse
+import sys
+import time
+import types
 
 
 def _stub_render_deps():
@@ -60,8 +63,10 @@ def main():
     v_fast = verts(m_fast)
 
     def chamfer(a, b, k=20000):
-        if a.shape[0] > k: a = a[torch.randperm(a.shape[0], device=a.device)[:k]]
-        if b.shape[0] > k: b = b[torch.randperm(b.shape[0], device=b.device)[:k]]
+        if a.shape[0] > k:
+            a = a[torch.randperm(a.shape[0], device=a.device)[:k]]
+        if b.shape[0] > k:
+            b = b[torch.randperm(b.shape[0], device=b.device)[:k]]
         return float((torch.cdist(a, b).min(1).values.mean()
                       + torch.cdist(b, a).min(1).values.mean()) / 2)
     cham = chamfer(v_stock, v_fast)