fix: model weights setup — export scripts, ORT compat, bootstrap
- Add export_models.py for OpenCLIP ViT-B/32 and YOLOv8n ONNX export - Fix ArgMax(13) ORT ARM64 incompatibility by passing eot_indices as a separate ONNX input (computed outside the graph in embed.py) - Use legacy TorchScript exporter (dynamo=False) for IR version 9 compat - Upgrade onnxruntime to 1.18.1 - Rewrite bootstrap_models.py with clear separation of auto-downloadable models (YuNet, SFace) vs manually-exported ones (OpenCLIP, YOLOv8n) - Wire bootstrap into worker CMD (runs before Celery) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
177
backend/app/services/vision/export_models.py
Normal file
177
backend/app/services/vision/export_models.py
Normal file
@@ -0,0 +1,177 @@
|
||||
"""
|
||||
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_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
|
||||
exported = Path("yolov8n.onnx")
|
||||
if exported.exists():
|
||||
exported.rename(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_yolov8n(models_dir)
|
||||
|
||||
logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user