fix: handle lazy/GGUF flow models in HiCacheModelPatch (issue #1)

pipeline.models[key] is None at patch time for GGUF Trellis2 pipelines
(Aero-Ex/ComfyUI-Trellis2-GGUF) that lazy-load the flow DiT inside the
sampler. nn.Module.__setattr__ only routes real Modules into _modules, so
self.inner = None became a plain __dict__ attribute that nn.Module.__getattr__
never searches, and the fallback raised 'object has no attribute inner'.

- store inner via _set_inner: real Modules stay registered submodules, None/
  lazy placeholders are kept as a plain attribute; __getattr__ resolves both
- clear errors when a compute step / attribute access hits an unloaded model
- bind_inner(model) to attach the real DiT once it materializes
- _LazyPatchDict binds the model transparently when the sampler assigns it
  into pipeline.models[key] (TRELLIS repos), so the cache survives lazy load

hicache_patch.py

@@ -114,12 +114,16 @@ class HiCacheModelPatch(torch.nn.Module):
     submodule, so ``.to()`` / ``.half()`` / state-dict access recurse into it).
     """
 
-    def __init__(self, model: torch.nn.Module, *, method: str = "hermite",
+    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:
         validate_config(method, interval, warmup_steps, max_order, sigma, dmd_history)
         super().__init__()
-        self.inner = model
+        # ``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
+        # attribute so __getattr__ never dead-ends. See _set_inner / bind_inner.
+        self._set_inner(model)
         self._hicache_is_patch = True  # marker for apply/remove
         self.method = method
         self.interval = int(interval)
@@ -134,14 +138,54 @@ class HiCacheModelPatch(torch.nn.Module):
         self.computed_steps = 0
         self.skipped_steps = 0
 
-    # -- attribute passthrough -------------------------------------------------
+    # -- 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
+        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
+        ``nn.Module.__getattr__`` never searches ``__dict__`` -- so the old
+        ``__getattr__`` fallback raised the misleading
+        ``'HiCacheModelPatch' object has no attribute 'inner'`` for lazy pipelines.
+        """
+        for d in (self.__dict__, self._parameters, self._buffers, self._modules):
+            d.pop("inner", None)
+        if isinstance(model, torch.nn.Module):
+            self._modules["inner"] = model
+        else:
+            object.__setattr__(self, "inner", model)
+
+    def _inner(self) -> Any:
+        """The wrapped model from wherever it lives (submodule or plain attr)."""
+        mod = self.__dict__.get("_modules")
+        if mod is not None and "inner" in mod:
+            return mod["inner"]
+        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,
+        registering it as a submodule. Lets lazy / GGUF pipelines that materialize
+        the model after patching still route through the cache."""
+        self._set_inner(model)
+        return self
+
     def __getattr__(self, name: str):
         try:
             return super().__getattr__(name)
         except AttributeError:
-            # fall through to the wrapped DiT (guidance_embed, config, dtype, ...)
-            inner = super().__getattr__("inner")
-            return getattr(inner, name)
+            pass
+        # fall through to the wrapped DiT (guidance_embed, config, dtype, ...)
+        inner = self._inner()
+        if inner is None:
+            raise AttributeError(
+                f"{type(self).__name__!r} has no attribute {name!r}: the wrapped DiT "
+                f"is not loaded yet (inner is None). Lazy / GGUF pipelines materialize "
+                f"the model inside the sampler; call bind_inner(model) once it exists."
+            )
+        return getattr(inner, name)
 
     # -- state handling ----------------------------------------------------------
     def _fresh_state(self) -> Dict[str, Any]:
@@ -205,7 +249,14 @@ class HiCacheModelPatch(torch.nn.Module):
             self.skipped_steps += 1
             return out
 
-        out = self.inner(latent_model_input, timestep, *args, **kwargs)
+        inner = self._inner()
+        if inner is None:
+            raise RuntimeError(
+                "HiCacheModelPatch: a compute step was reached but the wrapped DiT "
+                "is not loaded (inner is None). For lazy / GGUF pipelines, apply "
+                "HiCache after the model is materialized, or call bind_inner(model)."
+            )
+        out = inner(latent_model_input, timestep, *args, **kwargs)
         anchor = out.detach()
         hicache_update_derivatives(state, anchor)
         if self.method in ("dmd", "auto"):

tests/test_patch.py

@@ -341,6 +341,47 @@ def test_attribute_and_device_passthrough():
     assert pipe.model.inner is dit
 
 
+def test_wrap_none_does_not_crash_on_attr_access():
+    """Regression: lazy / GGUF pipelines may hand the wrapper a not-yet-loaded
+    (None) model. Wrapping None must not raise, and an unknown-attribute lookup
+    must give a clear 'not loaded yet' error, NOT the misleading
+    'HiCacheModelPatch object has no attribute inner'."""
+    patch = hp.HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2)
+    assert patch.inner is None
+    with pytest.raises(AttributeError) as ei:
+        _ = patch.guidance_embed          # forwarded lookup on a None inner
+    msg = str(ei.value)
+    assert "not loaded yet" in msg and "guidance_embed" in msg
+    with pytest.raises(RuntimeError):     # a compute step on an unbound patch
+        patch(torch.zeros(2, 8, 4), torch.tensor([0.0, 0.0]), cond=None)
+
+
+def test_bind_inner_materializes_lazy_model():
+    """After the real DiT loads, bind_inner attaches it (as a submodule) and the
+    patch runs/skip-forecasts through the denoise loop."""
+    patch = hp.HiCacheModelPatch(None, method="hermite", interval=3, warmup_steps=2)
+    dit = MockDiT()
+    assert patch.bind_inner(dit) is patch
+    assert patch.inner is dit
+    assert "inner" in patch._modules              # registered submodule
+    assert patch.guidance_embed is False          # passthrough works now
+    steps = 20
+    for t in torch.linspace(0, 1, steps):
+        patch(torch.zeros(2, 8, 4), t.expand(2), cond=None)
+    assert patch.skipped_steps > 0
+    assert dit.calls == patch.computed_steps
+
+
+def test_eager_inner_still_registered_as_submodule():
+    """Eager path unbroken: a real nn.Module inner stays a registered submodule so
+    state_dict / parameters / device moves recurse into it."""
+    dit = MockDiT()
+    patch = hp.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())
+
+
 def test_remove_restores_original_model():
     dit = MockDiT()
     pipe = MockPipeline(dit)