saving kv-cache

.gitignore

@@ -1 +1,70 @@
-demo_cache/
+demo_cache/
+
+
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache

stream3r/stream_session.py

@@ -1,3 +1,6 @@
+import os
+from typing import Any, Dict, Optional
+
 import torch
 from stream3r.models.stream3r import STream3R
 
@@ -68,7 +71,7 @@ class StreamSession:
 
     def _get_cache(self):
         return self.aggregator_kv_cache_list, self.camera_head_kv_cache_list
-    
+
     def get_all_predictions(self):
         return self.predictions
     
@@ -83,6 +86,96 @@ class StreamSession:
         self._clear_predictions()
         self._clear_cache()
 
+    @staticmethod
+    def _detach_to_cpu(cache_like):
+        if isinstance(cache_like, torch.Tensor):
+            return cache_like.detach().cpu()
+        if isinstance(cache_like, list):
+            return [StreamSession._detach_to_cpu(elem) for elem in cache_like]
+        if isinstance(cache_like, tuple):
+            return tuple(StreamSession._detach_to_cpu(elem) for elem in cache_like)
+        return cache_like
+
+    @staticmethod
+    def _to_device(cache_like, device: torch.device):
+        if isinstance(cache_like, torch.Tensor):
+            return cache_like.to(device)
+        if isinstance(cache_like, list):
+            return [StreamSession._to_device(elem, device) for elem in cache_like]
+        if isinstance(cache_like, tuple):
+            return tuple(StreamSession._to_device(elem, device) for elem in cache_like)
+        return cache_like
+
+    def _default_device(self) -> torch.device:
+        try:
+            return next(self.model.parameters()).device
+        except StopIteration:
+            return torch.device("cpu")
+
+    def save_cache(self, file_path: str) -> None:
+        aggregator_cache, camera_cache = self._get_cache()
+
+        payload: Dict[str, Any] = {
+            "metadata": {
+                "mode": self.mode,
+                "aggregator_depth": self.aggregator_kv_cache_depth,
+                "camera_head_depth": self.camera_head_kv_cache_depth,
+                "camera_head_iterations": self.camera_head_iterations,
+                "patch_size": getattr(self.model.aggregator, "patch_size", None),
+                "patch_start_idx": getattr(self.model.aggregator, "patch_start_idx", None),
+            },
+            "aggregator_cache": self._detach_to_cpu(aggregator_cache),
+            "camera_cache": self._detach_to_cpu(camera_cache),
+            "predictions": {k: v.detach().cpu() for k, v in self.predictions.items()},
+        }
+
+        dir_name = os.path.dirname(file_path)
+        if dir_name:
+            os.makedirs(dir_name, exist_ok=True)
+
+        torch.save(payload, file_path)
+
+    def load_cache(self, file_path: str, *, device: Optional[torch.device] = None, strict: bool = True) -> None:
+        if device is None:
+            device = self._default_device()
+
+        payload = torch.load(file_path, map_location="cpu")
+
+        metadata: Dict[str, Any] = payload.get("metadata", {})
+
+        expected_metadata = {
+            "mode": self.mode,
+            "aggregator_depth": self.aggregator_kv_cache_depth,
+            "camera_head_depth": self.camera_head_kv_cache_depth,
+            "camera_head_iterations": self.camera_head_iterations,
+        }
+
+        for key, expected_value in expected_metadata.items():
+            actual_value = metadata.get(key)
+            if strict and actual_value != expected_value:
+                raise ValueError(
+                    f"Loaded cache metadata mismatch for '{key}': expected {expected_value}, got {actual_value}"
+                )
+
+        if strict:
+            patch_size = getattr(self.model.aggregator, "patch_size", None)
+            patch_start_idx = getattr(self.model.aggregator, "patch_start_idx", None)
+            if metadata.get("patch_size") not in (None, patch_size):
+                raise ValueError(
+                    f"Loaded cache metadata mismatch for 'patch_size': expected {patch_size}, got {metadata.get('patch_size')}"
+                )
+            if metadata.get("patch_start_idx") not in (None, patch_start_idx):
+                raise ValueError(
+                    f"Loaded cache metadata mismatch for 'patch_start_idx': expected {patch_start_idx}, got {metadata.get('patch_start_idx')}"
+                )
+
+        self.aggregator_kv_cache_list = self._to_device(payload.get("aggregator_cache", []), device)
+        self.camera_head_kv_cache_list = self._to_device(payload.get("camera_cache", []), device)
+        self.predictions = {
+            k: v.to(device) if isinstance(v, torch.Tensor) else v
+            for k, v in payload.get("predictions", {}).items()
+        }
+
     def forward_stream(self, images):
         aggregator_kv_cache_list, camera_head_kv_cache_list = self._get_cache()
 

tests/test_stream_session_cache.py

@@ -0,0 +1,106 @@
+import os
+import tempfile
+import unittest
+
+try:
+    import torch
+except ImportError:  # pragma: no cover - environment without torch
+    torch = None
+
+if torch is not None:
+    from stream3r.stream_session import StreamSession
+
+
+if torch is None:
+
+    class StreamSessionCacheTest(unittest.TestCase):  # pragma: no cover - requires torch
+        @unittest.skip("PyTorch is required for this test")
+        def test_requires_torch(self):
+            pass
+
+else:
+
+    class _DummyAggregator:
+        depth = 2
+        patch_size = 4
+        patch_start_idx = 3
+
+
+    class _DummyCameraHead:
+        trunk_depth = 3
+
+
+    class _DummyModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.aggregator = _DummyAggregator()
+            self.camera_head = _DummyCameraHead()
+            self.register_parameter("_dummy_param", torch.nn.Parameter(torch.zeros(1)))
+
+        def forward(self, *args, **kwargs):
+            raise NotImplementedError
+
+
+    class StreamSessionCacheTest(unittest.TestCase):
+        def _make_populated_session(self) -> StreamSession:
+            model = _DummyModel()
+            session = StreamSession(model, mode="causal")
+
+            aggregator_cache = [
+                [torch.randn(1, 2, 3), torch.randn(1, 2, 3)]
+                for _ in range(session.aggregator_kv_cache_depth)
+            ]
+
+            camera_cache = []
+            for _ in range(session.camera_head_iterations):
+                iter_cache = []
+                for _ in range(session.camera_head_kv_cache_depth):
+                    iter_cache.append([torch.randn(1, 4, 5), torch.randn(1, 4, 5)])
+                camera_cache.append(iter_cache)
+
+            session.aggregator_kv_cache_list = aggregator_cache
+            session.camera_head_kv_cache_list = camera_cache
+            session.predictions = {
+                "depth": torch.randn(1, 2, 3, 3),
+                "pose_enc": torch.randn(1, 2, 9),
+            }
+
+            return session
+
+        def test_round_trip_save_and_load(self):
+            session = self._make_populated_session()
+            model = session.model
+
+            with tempfile.TemporaryDirectory() as tmpdir:
+                file_path = os.path.join(tmpdir, "kv_cache.pt")
+                session.save_cache(file_path)
+
+                restored_session = StreamSession(model, mode="causal")
+                restored_session.load_cache(file_path)
+
+            for original_layer, restored_layer in zip(
+                session.aggregator_kv_cache_list, restored_session.aggregator_kv_cache_list
+            ):
+                for original_tensor, restored_tensor in zip(original_layer, restored_layer):
+                    if original_tensor is None:
+                        self.assertIsNone(restored_tensor)
+                    else:
+                        self.assertTrue(torch.equal(original_tensor, restored_tensor))
+
+            for original_iter, restored_iter in zip(
+                session.camera_head_kv_cache_list, restored_session.camera_head_kv_cache_list
+            ):
+                for original_layer, restored_layer in zip(original_iter, restored_iter):
+                    for original_tensor, restored_tensor in zip(original_layer, restored_layer):
+                        if original_tensor is None:
+                            self.assertIsNone(restored_tensor)
+                        else:
+                            self.assertTrue(torch.equal(original_tensor, restored_tensor))
+
+            for key, original_tensor in session.predictions.items():
+                restored_tensor = restored_session.predictions[key]
+                self.assertTrue(torch.equal(original_tensor, restored_tensor))
+
+
+if __name__ == "__main__":  # pragma: no cover - manual execution
+    unittest.main()