feat: upgrade to SigLIP2 ViT-B/16 for semantic search
Replace OpenCLIP ViT-B/32 (512-d, ~78% recall) with SigLIP2 ViT-B/16 (768-d, ~84% recall) as the default embedding model for significantly better image-text retrieval quality. - New SigLIP2Embedder class with 384px input and SigLIP normalization - ONNX export pipeline for SigLIP2 visual + textual encoders - Migration 0010: resize embeddings.vector from 512 to 768 dimensions - Config-driven model selection: "siglip2_vitb16" (default) or "openclip_vitb32" (legacy) — both models can coexist - Content classifier follows the configured embedder family - Existing embeddings cleared on migration; vision backfill regenerates Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -120,6 +120,76 @@ def export_openclip(models_dir: Path):
|
||||
logger.info("Textual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
|
||||
def export_siglip2(models_dir: Path):
|
||||
"""Export SigLIP2 ViT-B/16 to two ONNX files (visual + textual)."""
|
||||
import torch
|
||||
import open_clip
|
||||
|
||||
out_dir = models_dir / "embed_siglip2"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
visual_path = out_dir / "visual.onnx"
|
||||
textual_path = out_dir / "textual.onnx"
|
||||
|
||||
if visual_path.exists() and textual_path.exists():
|
||||
logger.info("SigLIP2 ONNX files already exist, skipping export")
|
||||
return
|
||||
|
||||
logger.info("Loading SigLIP2 ViT-B-16-SigLIP2 webli...")
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
"ViT-B-16-SigLIP2", pretrained="webli"
|
||||
)
|
||||
model.eval()
|
||||
|
||||
export_kwargs = dict(opset_version=14, dynamo=False)
|
||||
|
||||
# ── Visual encoder ────────────────────────────────────────────────
|
||||
if not visual_path.exists():
|
||||
logger.info("Exporting SigLIP2 visual encoder → %s", visual_path)
|
||||
dummy_image = torch.randn(1, 3, 384, 384)
|
||||
torch.onnx.export(
|
||||
model.visual,
|
||||
dummy_image,
|
||||
str(visual_path),
|
||||
input_names=["image"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"image": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = visual_path.stat().st_size / 1e6
|
||||
logger.info("SigLIP2 visual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
# ── Textual encoder ───────────────────────────────────────────────
|
||||
if not textual_path.exists():
|
||||
logger.info("Exporting SigLIP2 textual encoder → %s", textual_path)
|
||||
tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP2")
|
||||
dummy_text = tokenizer(["a photo"]).to(torch.int64)
|
||||
|
||||
class SigLIP2TextEncoder(torch.nn.Module):
|
||||
"""Wrap the SigLIP2 text transformer for ONNX export."""
|
||||
def __init__(self, clip_model):
|
||||
super().__init__()
|
||||
self.text = clip_model.text
|
||||
|
||||
def forward(self, text):
|
||||
return self.text(text)
|
||||
|
||||
text_enc = SigLIP2TextEncoder(model)
|
||||
text_enc.eval()
|
||||
|
||||
torch.onnx.export(
|
||||
text_enc,
|
||||
dummy_text,
|
||||
str(textual_path),
|
||||
input_names=["text"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"text": {0: "batch"}},
|
||||
**export_kwargs,
|
||||
)
|
||||
size_mb = textual_path.stat().st_size / 1e6
|
||||
logger.info("SigLIP2 textual encoder exported (%.1f MB)", size_mb)
|
||||
|
||||
|
||||
def export_yolov8n(models_dir: Path):
|
||||
"""Export YOLOv8n to ONNX."""
|
||||
out_dir = models_dir / "detect"
|
||||
@@ -174,6 +244,7 @@ def main():
|
||||
logger.info("Exporting models to %s", models_dir)
|
||||
|
||||
export_openclip(models_dir)
|
||||
export_siglip2(models_dir)
|
||||
export_yolov8n(models_dir)
|
||||
|
||||
logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")
|
||||
|
||||
Reference in New Issue
Block a user