File size: 5,425 Bytes
76d7a6e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """Memory-mapped frozen semantic features used by Dense UV Parser."""
import json
from pathlib import Path
import numpy as np
import torch
SIGLIP_CACHE_VERSION = 2
SUPPORTED_SIGLIP_CACHE_VERSIONS = (1, 2)
class SigLIPGlobalCache:
def __init__(
self,
cache_dir,
expected_views=None,
expected_model=None,
expected_data_dir=None,
require_spatial=False,
):
self.cache_dir = Path(cache_dir)
metadata_path = self.cache_dir / "metadata.json"
embeddings_path = self.cache_dir / "embeddings.npy"
if not metadata_path.is_file() or not embeddings_path.is_file():
raise FileNotFoundError(
f"Incomplete SigLIP cache in {self.cache_dir}; expected metadata.json "
"and embeddings.npy."
)
self.metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
if self.metadata.get("version") not in SUPPORTED_SIGLIP_CACHE_VERSIONS:
raise ValueError(
f"Unsupported SigLIP cache version {self.metadata.get('version')!r}."
)
if expected_views is not None and self.metadata.get("views") != list(
expected_views
):
raise ValueError(
f"SigLIP cache views={self.metadata.get('views')!r}, "
f"expected {list(expected_views)!r}."
)
if (
expected_model is not None
and self.metadata.get("siglip_model") != expected_model
):
raise ValueError(
f"SigLIP cache model={self.metadata.get('siglip_model')!r}, "
f"expected {expected_model!r}."
)
if expected_data_dir is not None:
cached_data_dir = self.metadata.get("data_dir")
requested_data_dir = Path(expected_data_dir).resolve()
if cached_data_dir and Path(cached_data_dir).resolve() != requested_data_dir:
raise ValueError(
f"SigLIP cache data_dir={str(Path(cached_data_dir).resolve())!r}, "
f"expected {str(requested_data_dir)!r}."
)
filenames = self.metadata.get("filenames")
if not isinstance(filenames, list) or len(set(filenames)) != len(filenames):
raise ValueError("SigLIP cache filenames must be a unique list.")
self.filename_to_index = {
filename: index for index, filename in enumerate(filenames)
}
self.embeddings = np.load(embeddings_path, mmap_mode="r")
expected_shape = (
len(filenames),
len(self.metadata["views"]),
int(self.metadata["feature_dim"]),
)
if self.embeddings.shape != expected_shape:
raise ValueError(
f"SigLIP cache shape={self.embeddings.shape}, expected {expected_shape}."
)
spatial_metadata = (
self.metadata.get("spatial_feature_dim"),
self.metadata.get("spatial_height"),
self.metadata.get("spatial_width"),
)
self.has_spatial = all(value is not None for value in spatial_metadata)
self.spatial_embeddings = None
if self.has_spatial:
spatial_path = self.cache_dir / "spatial_embeddings.npy"
if not spatial_path.is_file():
raise FileNotFoundError(
"SigLIP cache metadata declares spatial features but "
f"{spatial_path} is missing."
)
self.spatial_embeddings = np.load(spatial_path, mmap_mode="r")
expected_spatial_shape = (
len(filenames),
len(self.metadata["views"]),
int(self.metadata["spatial_feature_dim"]),
int(self.metadata["spatial_height"]),
int(self.metadata["spatial_width"]),
)
if self.spatial_embeddings.shape != expected_spatial_shape:
raise ValueError(
"SigLIP spatial cache shape="
f"{self.spatial_embeddings.shape}, expected "
f"{expected_spatial_shape}."
)
elif any(value is not None for value in spatial_metadata):
raise ValueError("SigLIP spatial cache metadata is incomplete.")
if require_spatial and not self.has_spatial:
raise ValueError(
f"SigLIP cache {self.cache_dir} contains only global features; "
"rebuild it with spatial caching enabled."
)
def get(self, filename):
try:
index = self.filename_to_index[filename]
except KeyError as error:
raise KeyError(f"{filename} is missing from the SigLIP cache.") from error
return torch.from_numpy(
np.array(self.embeddings[index], dtype=np.float32, copy=True)
)
def get_spatial(self, filename):
if not self.has_spatial:
raise ValueError("This SigLIP cache has no spatial features.")
try:
index = self.filename_to_index[filename]
except KeyError as error:
raise KeyError(f"{filename} is missing from the SigLIP cache.") from error
return torch.from_numpy(
np.array(
self.spatial_embeddings[index],
dtype=np.float16,
copy=True,
)
)
|