mule-image/backend/app/services/vision/export_models.py

"""
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. 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_yolov8n(models_dir)

    logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")


if __name__ == "__main__":
    main()