From 3bb7b78557be67b0dc1e6bcabd5b95b31ff42e1b Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Tue, 25 Aug 2026 13:54:01 +0300 Subject: [PATCH] perf(studio): keep the key's tensors in a binary sidecar, not as base64 in its JSON MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One conditioning delta is 586x3584 floats — several megabytes of base64 inside a document that is parsed and rewritten IN FULL on every generation. That is the cost that grows with everything ever learned, and it is paid whether or not any of those tensors are used. Tensors now go to .json.tensors.pt and the JSON keeps a named reference whose id is the tensor's path in the tree ("state.global.liked_dir.direction"), so the key stays readable and a reader can see what a reference points at. No migration: a key written before this carries inline blobs, is read exactly as it stands, and moves to the sidecar on its next save. If the sidecar cannot be written the tensors stay inline — a slow key is fine, a JSON pointing at a sidecar that does not exist is not. A reference the sidecar has lost reads as no tensor, so a direction goes unapplied rather than breaking the run. Deletion already sweeps it: delete_refinement_key globs every .* sidecar for exactly this reason. Session Reset removes it explicitly. Co-Authored-By: Claude Opus 5 --- conditioning.py | 120 ++++++++++++++++++++++++++++++++++- tests/test_tensor_sidecar.py | 81 +++++++++++++++++++++++ 2 files changed, 200 insertions(+), 1 deletion(-) create mode 100644 tests/test_tensor_sidecar.py diff --git a/conditioning.py b/conditioning.py index faf95337..8282f4cf 100644 --- a/conditioning.py +++ b/conditioning.py @@ -252,13 +252,72 @@ def tensor_to_serializable(t: torch.Tensor) -> dict: "dtype": str(arr.dtype) } +# A tensor inside the refinement key is one of three things: +# {"data","shape","dtype"} an inline base64 blob — the original format, still read +# {"__tensor__": Tensor} loaded from the sidecar, in memory only +# {"__tensor_ref__": id} what an externalised tensor looks like ON DISK +# The blob form is what makes a key slow: it is base64 inside a JSON that is fully parsed and +# fully rewritten every run, and one conditioning delta is several megabytes of it. +TENSOR_SIDECAR_SUFFIX = ".tensors.pt" + + def serializable_to_tensor(d: dict) -> torch.Tensor: + held = d.get("__tensor__") if isinstance(d, dict) else None + if isinstance(held, torch.Tensor): + tensor = held.to(dtype=torch.float32) + return tensor.cuda() if torch.cuda.is_available() else tensor arr = np.frombuffer(base64.b64decode(d["data"]), dtype=d["dtype"]).reshape(d["shape"]).copy() tensor = torch.from_numpy(arr).to(dtype=torch.float32) if torch.cuda.is_available(): tensor = tensor.cuda() return tensor +def _is_tensor_blob(value): + return isinstance(value, dict) and "data" in value and "shape" in value and "dtype" in value + + +def externalize_tensors(value, store, path="state"): + """Replace every tensor in a state tree with a reference, collecting them into `store`. + + Ids are the tensor's PATH in the tree ("state.global.liked_dir.direction"), so the JSON + stays readable and diffable and a reader can see what a reference points at. Returns the + tree with references in place of blobs; the caller writes `store` beside the JSON. + """ + if isinstance(value, dict): + held = value.get("__tensor__") + if isinstance(held, torch.Tensor): + store[path] = held.detach().cpu() + return {"__tensor_ref__": path, "shape": list(held.shape), "dtype": str(held.dtype)} + if _is_tensor_blob(value): + arr = np.frombuffer(base64.b64decode(value["data"]), + dtype=value["dtype"]).reshape(value["shape"]).copy() + store[path] = torch.from_numpy(arr) + return {"__tensor_ref__": path, "shape": list(value["shape"]), + "dtype": str(value["dtype"])} + return {k: externalize_tensors(v, store, f"{path}.{k}") for k, v in value.items()} + if isinstance(value, list): + return [externalize_tensors(v, store, f"{path}[{i}]") for i, v in enumerate(value)] + return value + + +def internalize_tensors(value, store): + """Put the sidecar's tensors back behind their references, in memory. + + A reference whose tensor is missing from the sidecar is left as-is: consumers read + tensors through serializable_to_tensor, which returns nothing for it, and a missing + direction is a direction not applied rather than a broken run. + """ + if isinstance(value, dict): + ref = value.get("__tensor_ref__") + if isinstance(ref, str): + held = store.get(ref) + return {"__tensor__": held} if isinstance(held, torch.Tensor) else value + return {k: internalize_tensors(v, store) for k, v in value.items()} + if isinstance(value, list): + return [internalize_tensors(v, store) for v in value] + return value + + def _safe_float(value, fallback=0.0): try: return float(value) @@ -3934,6 +3993,14 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): clear_refinement_data(refinement_key) except Exception as e: print(f"[FunPackVideoRefinerV2] Enhancement cleanup failed: {e}") + # The tensor sidecar goes with it. Reset means fresh, and a sidecar left behind + # would be read back by the next load if the fresh state were never saved. + for _sidecar in (self._v2_tensor_sidecar_path(path),): + try: + if os.path.exists(_sidecar): + os.remove(_sidecar) + except OSError as error: + print(f"[FunPackVideoRefinerV2] tensor sidecar cleanup failed: {error}") vf_path = refinement_state_path(refinement_key, "value_fn", prefix="refine_v2", extension="pt") try: if os.path.exists(vf_path): @@ -3961,6 +4028,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): _started = time.perf_counter() with open(path, "r", encoding="utf-8") as file: data = json.load(file) + data = self._v2_load_state_tensors(path, data) self._v2_note_state_io("load", time.perf_counter() - _started, os.path.getsize(path), data) if not isinstance(data, dict) or int(data.get("version", 0)) != 2: @@ -4031,14 +4099,64 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): path = self._v2_state_path(refinement_key) os.makedirs(os.path.dirname(path), exist_ok=True) _started = time.perf_counter() + payload = self._v2_store_state_tensors(path, data) with open(path, "w", encoding="utf-8") as file: - json.dump(self._v2_json_safe(data), file, indent=2) + json.dump(self._v2_json_safe(payload), file, indent=2) try: self._v2_note_state_io("save", time.perf_counter() - _started, os.path.getsize(path), data) except OSError: pass + @staticmethod + def _v2_tensor_sidecar_path(state_path): + return str(state_path) + TENSOR_SIDECAR_SUFFIX + + def _v2_load_state_tensors(self, state_path, data): + """Resolve the key's tensor references against its sidecar, if it has one. + + A key written before the sidecar existed carries inline blobs and is returned + untouched — serializable_to_tensor still reads those, and the next save moves them + out. No migration step, no version bump. + """ + sidecar = self._v2_tensor_sidecar_path(state_path) + if not os.path.exists(sidecar): + return data + try: + store = torch.load(sidecar, map_location="cpu", weights_only=True) + except Exception as error: # noqa: BLE001 + _log.failed("FunPackVideoRefinerV2", "refinement key tensor sidecar", error, + "the directions and deltas it holds are not applied this run; the " + "rest of the key is used normally and the next save rewrites it") + return data + return internalize_tensors(data, store if isinstance(store, dict) else {}) + + def _v2_store_state_tensors(self, state_path, data): + """Write the key's tensors to a binary sidecar; return the JSON-bound tree. + + Conditioning deltas and learned directions are megabytes each as base64, inside a + document that is parsed and rewritten in full on every run. Held as tensors in a .pt + beside the key, the JSON keeps only a named reference and stays readable. + + If the sidecar cannot be written, the tensors stay INLINE — an old-format key is + slow, but a JSON pointing at a sidecar that does not exist has lost them. + """ + store = {} + try: + payload = externalize_tensors(data, store, path="state") + if not store: + return payload + sidecar = self._v2_tensor_sidecar_path(state_path) + tmp = sidecar + ".tmp" + torch.save(store, tmp) + os.replace(tmp, sidecar) + return payload + except Exception as error: # noqa: BLE001 + _log.failed("FunPackVideoRefinerV2", "refinement key tensor sidecar write", error, + "the tensors are written inline in the JSON instead — the key still " + "works and stays correct, it is just the slower format") + return data + # The collections that can grow without bound, and are therefore the ones worth naming # when a key gets slow. Each is re-parsed and rewritten in full on every run. _V2_STATE_COLLECTIONS = ("phrase_memory", "conditioning_deltas", "concept_delta_memory", diff --git a/tests/test_tensor_sidecar.py b/tests/test_tensor_sidecar.py new file mode 100644 index 00000000..33900c7d --- /dev/null +++ b/tests/test_tensor_sidecar.py @@ -0,0 +1,81 @@ +"""Conditioning tensors live in a binary sidecar, not as base64 inside the key's JSON. + +One conditioning delta is 586x3584 floats — several megabytes of base64 in a document that +is parsed and rewritten in full on every generation. That is what makes a well-used key get +slower run after run. +""" +import json +import os +import sys + +import pytest +import torch + +sys.path.insert(0, ".") + + +@pytest.fixture +def C(): + import conditioning + return conditioning + + +def test_a_tensor_survives_the_round_trip(C): + tree = {"global": {"liked_dir": {"direction": C.tensor_to_serializable(torch.arange(6.0))}}} + store = {} + out = C.externalize_tensors(tree, store) + assert "state.global.liked_dir.direction" in store + back = C.internalize_tensors(json.loads(json.dumps(out)), store) + got = C.serializable_to_tensor(back["global"]["liked_dir"]["direction"]) + assert torch.equal(got.cpu().float(), torch.arange(6.0)) + + +def test_the_json_keeps_a_readable_reference(C): + tree = {"global": {"liked_dir": {"direction": C.tensor_to_serializable(torch.ones(4))}}} + out = C.externalize_tensors(tree, {}) + ref = out["global"]["liked_dir"]["direction"] + assert ref["__tensor_ref__"] == "state.global.liked_dir.direction" + assert "data" not in ref # the base64 is gone + assert ref["shape"] == [4] # still says what it points at + + +def test_no_tensors_means_no_sidecar(C): + store = {} + C.externalize_tensors({"global": {"phrase_memory": {"a": {"score": 1.0}}}}, store) + assert store == {} + + +def test_a_missing_tensor_is_not_a_broken_run(C): + """A reference the sidecar has lost reads as no tensor — a direction not applied.""" + tree = {"d": {"__tensor_ref__": "state.d", "shape": [4], "dtype": "float32"}} + back = C.internalize_tensors(tree, {}) + assert back["d"] == tree["d"] + with pytest.raises(Exception): + C.serializable_to_tensor(back["d"]) + + +def test_an_old_key_with_inline_blobs_still_reads(C): + """No migration step: a key written before the sidecar existed is read as it stands.""" + blob = C.tensor_to_serializable(torch.arange(3.0)) + assert torch.equal(C.serializable_to_tensor(blob).cpu().float(), torch.arange(3.0)) + + +def test_lists_are_walked_too(C): + tree = {"anchors": [{"embed": C.tensor_to_serializable(torch.zeros(2))}]} + store = {} + out = C.externalize_tensors(tree, store) + assert "state.anchors[0].embed" in store + assert out["anchors"][0]["embed"]["__tensor_ref__"] == "state.anchors[0].embed" + + +def test_the_state_written_to_disk_is_json(C, tmp_path): + """The whole point: what lands in the .json must contain no tensor payload.""" + from conditioning import FunPackVideoRefinerV2 + r = FunPackVideoRefinerV2() + path = str(tmp_path / "key.json") + tree = {"global": {"liked_dir": {"direction": C.tensor_to_serializable(torch.ones(1000))}}} + payload = r._v2_store_state_tensors(path, tree) + text = json.dumps(payload) + assert "__tensor_ref__" in text + assert len(text) < 500 + assert os.path.exists(path + C.TENSOR_SIDECAR_SUFFIX)