Spaces:
Sleeping
Sleeping
File size: 3,672 Bytes
e510bac | 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 | """Offline export of the SigLIP2 text embeddings for RoofSight's fixed prompts.
Run this ONCE on a machine with RAM headroom (~14 GB peak). It loads the full
`siglip2-g` adaptor — which instantiates the entire google/siglip2-giant model
just to keep its text tower — encodes the fixed prompt set, and writes the
resulting (num_prompts, D) tensor to disk.
The Space then loads with the `siglip2` adaptor factory removed from RADIO's
registry, so the text tower is never built there. See radio_backbone.load_model.
Also captures a reference feature map so the stripped runtime path can be
verified as numerically identical:
python export_text_embeddings.py # export + reference
python verify_stripped_adaptor.py # compare
"""
import os
import sys
import numpy as np
import torch
from PIL import Image
from radio_backbone import ROOF_PROMPTS, NON_ROOF_PROMPTS, PATCH_SIZE
OUT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "text_embeddings.pt")
REF_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_ref_features.pt")
REF_IMAGE = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"_run_output", "1310_susan_circle", "original.png",
)
# CPU + fp32 throughout: the reference must be deterministic so the parity
# check measures the adaptor swap, not autocast/GPU nondeterminism.
DEVICE = "cpu"
def main():
all_labels = ROOF_PROMPTS + NON_ROOF_PROMPTS
print(f"Loading FULL c-radio_v4-h + siglip2-g adaptor on {DEVICE} (this is the heavy path)...")
model = torch.hub.load(
"NVlabs/RADIO",
"radio_model",
version="c-radio_v4-h",
adaptor_names=["siglip2-g"],
vitdet_window_size=8,
progress=True,
skip_validation=True,
trust_repo=True,
)
model.eval()
model.to(DEVICE)
adaptor = model.adaptors["siglip2-g"]
assert hasattr(adaptor, "encode_text"), (
"adaptor has no encode_text — the siglip2 factory was not used. "
"This script must run WITHOUT the registry patch."
)
print(f"Encoding {len(all_labels)} prompts...")
text_input = adaptor.tokenizer(all_labels)
text_input = {k: v.to(DEVICE) for k, v in text_input.items()}
with torch.no_grad():
embeddings = adaptor.encode_text(text_input, normalize=True).float().cpu().clone()
print(f" embeddings: {tuple(embeddings.shape)} dtype={embeddings.dtype}")
torch.save({"prompts": all_labels, "embeddings": embeddings}, OUT_PATH)
print(f" wrote {OUT_PATH} ({os.path.getsize(OUT_PATH) / 1024:.1f} KB)")
# --- Reference feature map, for the parity check ---
if not os.path.exists(REF_IMAGE):
print(f"\nNo reference image at {REF_IMAGE}; skipping reference capture.")
return
print(f"\nCapturing reference features from {REF_IMAGE}...")
image = np.array(Image.open(REF_IMAGE).convert("RGB"))
x = torch.from_numpy(image).permute(2, 0, 1).float().div_(255.0).unsqueeze(0)
snapped = model.get_nearest_supported_resolution(*x.shape[-2:])
if snapped != x.shape[-2:]:
x = torch.nn.functional.interpolate(x, snapped, mode="bilinear", align_corners=False)
with torch.no_grad():
out = model(x.to(DEVICE))
_, features = out["siglip2-g"]
torch.save(
{
"features": features.float().cpu().clone(),
"snapped": tuple(snapped),
"image_shape": image.shape,
},
REF_PATH,
)
print(f" features: {tuple(features.shape)}")
print(f" wrote {REF_PATH}")
print("\nDone. Now run: python verify_stripped_adaptor.py")
if __name__ == "__main__":
sys.exit(main())
|