""" 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-SigLIP-384 webli...") model, _, preprocess = open_clip.create_model_and_transforms( "ViT-B-16-SigLIP-384", 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-SigLIP-384") 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()