refactor: strip AI pipeline to binary photo/other classifier
Drops face recognition, OCR, object detection, and semantic embeddings. The sole remaining vision task is a CLIP-based binary classifier (photography vs other); photos in "other" get needs_review=true so screenshots, documents, memes and scans can be triaged from a new filter pill in the UI. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,253 +1,61 @@
|
||||
"""
|
||||
Export / download all vision model weights to ONNX format.
|
||||
Export the OpenCLIP ViT-B/32 visual encoder to ONNX.
|
||||
|
||||
Run ONCE on any machine with Python + pip (doesn't need GPU):
|
||||
Run once on any machine with Python + pip (no GPU needed):
|
||||
|
||||
pip install open-clip-torch ultralytics onnx
|
||||
pip install open-clip-torch 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.
|
||||
Produces:
|
||||
embed/visual.onnx (~350 MB)
|
||||
"""
|
||||
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)."""
|
||||
def export_openclip_visual(models_dir: Path):
|
||||
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")
|
||||
if visual_path.exists():
|
||||
logger.info("OpenCLIP visual.onnx already exists, skipping export")
|
||||
return
|
||||
|
||||
logger.info("Loading OpenCLIP ViT-B-32 laion2b_s34b_b79k...")
|
||||
model, _, preprocess = open_clip.create_model_and_transforms(
|
||||
model, _, _ = 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"
|
||||
logger.info("Exporting visual encoder → %s", visual_path)
|
||||
dummy = torch.randn(1, 3, 224, 224)
|
||||
torch.onnx.export(
|
||||
model.visual,
|
||||
dummy,
|
||||
str(visual_path),
|
||||
input_names=["image"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={"image": {0: "batch"}},
|
||||
opset_version=14,
|
||||
dynamo=False,
|
||||
)
|
||||
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)
|
||||
logger.info("Visual encoder exported (%.1f MB)", visual_path.stat().st_size / 1e6)
|
||||
|
||||
|
||||
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)",
|
||||
)
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--models-dir", type=Path, default=Path("/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.")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
args.models_dir.mkdir(parents=True, exist_ok=True)
|
||||
export_openclip_visual(args.models_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user