dtoro
94c07b1d0d
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>
255 lines
9.1 KiB
Python
255 lines
9.1 KiB
Python
"""
|
|
Export / download all vision model weights to ONNX format.
|
|
|
|
Run ONCE on any machine with Python + pip (doesn't need GPU):
|
|
|
|
pip install open-clip-torch ultralytics onnx
|
|
python -m app.services.vision.export_models [--models-dir /data/models]
|
|
|
|
This produces:
|
|
embed/visual.onnx (~350 MB)
|
|
embed/textual.onnx (~250 MB)
|
|
detect/yolov8n.onnx (~12 MB)
|
|
|
|
YuNet and SFace are downloaded by bootstrap_models.py at worker boot
|
|
(Apache 2.0, lightweight, no export step needed).
|
|
|
|
After export, copy the /data/models directory into your Docker volume:
|
|
docker cp /data/models mulita-worker:/data/models
|
|
Or mount a host path in docker-compose.yml.
|
|
"""
|
|
import argparse
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def export_openclip(models_dir: Path):
|
|
"""Export OpenCLIP ViT-B/32 to two ONNX files (visual + textual)."""
|
|
import torch
|
|
import open_clip
|
|
|
|
out_dir = models_dir / "embed"
|
|
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("OpenCLIP ONNX files already exist, skipping export")
|
|
return
|
|
|
|
logger.info("Loading OpenCLIP ViT-B-32 laion2b_s34b_b79k...")
|
|
model, _, preprocess = open_clip.create_model_and_transforms(
|
|
"ViT-B-32", pretrained="laion2b_s34b_b79k"
|
|
)
|
|
model.eval()
|
|
|
|
# Use dynamo=False to get the legacy TorchScript exporter which
|
|
# produces IR version 9 (compatible with onnxruntime 1.17.x).
|
|
# The new torch.onnx.export default (dynamo=True) emits IR 10.
|
|
export_kwargs = dict(opset_version=14, dynamo=False)
|
|
|
|
# ── Visual encoder ────────────────────────────────────────────────
|
|
if not visual_path.exists():
|
|
logger.info("Exporting visual encoder → %s", visual_path)
|
|
dummy_image = torch.randn(1, 3, 224, 224)
|
|
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("Visual encoder exported (%.1f MB)", size_mb)
|
|
|
|
# ── Textual encoder ───────────────────────────────────────────────
|
|
if not textual_path.exists():
|
|
logger.info("Exporting textual encoder → %s", textual_path)
|
|
tokenizer = open_clip.get_tokenizer("ViT-B-32")
|
|
dummy_text = tokenizer(["a photo"]).to(torch.int64)
|
|
|
|
class TextEncoder(torch.nn.Module):
|
|
"""Wrap the CLIP text encoder to avoid argmax in the ONNX graph.
|
|
OpenCLIP uses argmax to find the EOT token position, but ORT
|
|
ARM64 doesn't support ArgMax(13). We pre-compute the EOT index
|
|
from the token sequence and pass it directly."""
|
|
def __init__(self, clip_model):
|
|
super().__init__()
|
|
self.transformer = clip_model.transformer
|
|
self.token_embedding = clip_model.token_embedding
|
|
self.positional_embedding = clip_model.positional_embedding
|
|
self.ln_final = clip_model.ln_final
|
|
self.text_projection = clip_model.text_projection
|
|
|
|
def forward(self, text, eot_indices):
|
|
x = self.token_embedding(text)
|
|
x = x + self.positional_embedding
|
|
x = x.permute(1, 0, 2) # NLD -> LND
|
|
x = self.transformer(x)
|
|
x = x.permute(1, 0, 2) # LND -> NLD
|
|
x = self.ln_final(x)
|
|
# Take the feature at the EOT token. The EOT index is
|
|
# passed in as a separate input (computed outside ONNX)
|
|
# to avoid ArgMax(13) which ORT ARM64 doesn't support.
|
|
x = x[torch.arange(x.shape[0]), eot_indices]
|
|
x = x @ self.text_projection
|
|
return x
|
|
|
|
text_enc = TextEncoder(model)
|
|
text_enc.eval()
|
|
|
|
# Compute EOT indices from dummy tokens (argmax of token ids)
|
|
dummy_eot = dummy_text.argmax(dim=-1)
|
|
|
|
torch.onnx.export(
|
|
text_enc,
|
|
(dummy_text, dummy_eot),
|
|
str(textual_path),
|
|
input_names=["text", "eot_indices"],
|
|
output_names=["embedding"],
|
|
dynamic_axes={"text": {0: "batch"}, "eot_indices": {0: "batch"}},
|
|
**export_kwargs,
|
|
)
|
|
size_mb = textual_path.stat().st_size / 1e6
|
|
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"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
onnx_path = out_dir / "yolov8n.onnx"
|
|
|
|
if onnx_path.exists():
|
|
logger.info("YOLOv8n ONNX already exists, skipping export")
|
|
return
|
|
|
|
logger.info("Exporting YOLOv8n → %s", onnx_path)
|
|
|
|
from ultralytics import YOLO
|
|
|
|
model = YOLO("yolov8n.pt")
|
|
model.export(format="onnx", imgsz=640, simplify=True)
|
|
|
|
# ultralytics exports to cwd as yolov8n.onnx — move to target. Use
|
|
# shutil.move rather than Path.rename so it works across filesystems
|
|
# (the cwd is typically /app inside the container, while the target
|
|
# /data/models is a separately-mounted volume — Path.rename raises
|
|
# "Invalid cross-device link" in that case).
|
|
import shutil
|
|
|
|
exported = Path("yolov8n.onnx")
|
|
if exported.exists():
|
|
shutil.move(str(exported), str(onnx_path))
|
|
|
|
size_mb = onnx_path.stat().st_size / 1e6
|
|
logger.info("YOLOv8n exported (%.1f MB)", size_mb)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Export vision model weights to ONNX")
|
|
parser.add_argument(
|
|
"--models-dir",
|
|
type=Path,
|
|
default=Path("/data/models"),
|
|
help="Directory to write model files (default: /data/models)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
)
|
|
|
|
models_dir = args.models_dir
|
|
models_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
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.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|