prompt cache: preserve frozen tensor dtypes across serialization

tests/test_advanced_prompt_embeds.py

@@ -0,0 +1,20 @@
+import torch
+
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+
+
+def test_frozen_dtype_keys_survive_save_load_and_dtype_conversion(tmp_path):
+    cache_path = tmp_path / "prompt.safetensors"
+    original = AdvancedPromptEmbeds(
+        text_embeds=[torch.ones(2, dtype=torch.float32)],
+        token_ids=[torch.tensor([1, 2], dtype=torch.int64)],
+    )
+    original.frozen_dtype_keys = ["token_ids"]
+
+    original.save(str(cache_path))
+    loaded = AdvancedPromptEmbeds.load(str(cache_path))
+    converted = loaded.to(dtype=torch.bfloat16)
+
+    assert loaded.frozen_dtype_keys == ["token_ids"]
+    assert converted.text_embeds[0].dtype == torch.bfloat16
+    assert converted.token_ids[0].dtype == torch.int64

toolkit/advanced_prompt_embeds.py

@@ -1,5 +1,7 @@
 import os
+import json
 import torch
+from safetensors import safe_open
 from safetensors.torch import load_file, save_file
 
 
@@ -136,7 +138,10 @@ def expand_to_batch(self, batch_size):
 
     def save(self, path):
         data = {}
-        metadata = {"class_name": self.__class__.__name__}
+        metadata = {
+            "class_name": self.__class__.__name__,
+            "frozen_dtype_keys": json.dumps(self._frozen_dtype_keys),
+        }
         for key, value in self._store.items():
             if len(value) != 1:
                 raise ValueError(
@@ -149,6 +154,8 @@ def save(self, path):
     @classmethod
     def load(cls, path=None):
         if path is not None:
+            with safe_open(path, framework="pt", device="cpu") as f:
+                metadata = f.metadata() or {}
             loaded = load_file(path)
         else:
             raise ValueError("Must provide a path")
@@ -157,7 +164,18 @@ def load(cls, path=None):
         for key in loaded.keys():
             data[key] = loaded[key]
 
-        return cls(**data)
+        prompt_embeds = cls(**data)
+        serialized_frozen_keys = metadata.get("frozen_dtype_keys")
+        if serialized_frozen_keys is not None:
+            try:
+                frozen_keys = json.loads(serialized_frozen_keys)
+            except (TypeError, json.JSONDecodeError):
+                frozen_keys = []
+            if isinstance(frozen_keys, list):
+                prompt_embeds.frozen_dtype_keys = [
+                    key for key in frozen_keys if isinstance(key, str)
+                ]
+        return prompt_embeds
 
     @classmethod
     def concat_prompt_embeds(