dtoro
574d71371f
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>
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
"""
|
|
Export the OpenCLIP ViT-B/32 visual encoder to ONNX.
|
|
|
|
Run once on any machine with Python + pip (no GPU needed):
|
|
|
|
pip install open-clip-torch onnx
|
|
python -m app.services.vision.export_models [--models-dir /data/models]
|
|
|
|
Produces:
|
|
embed/visual.onnx (~350 MB)
|
|
"""
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
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"
|
|
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, _, _ = open_clip.create_model_and_transforms(
|
|
"ViT-B-32", pretrained="laion2b_s34b_b79k"
|
|
)
|
|
model.eval()
|
|
|
|
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,
|
|
)
|
|
logger.info("Visual encoder exported (%.1f MB)", visual_path.stat().st_size / 1e6)
|
|
|
|
|
|
def main():
|
|
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")
|
|
args.models_dir.mkdir(parents=True, exist_ok=True)
|
|
export_openclip_visual(args.models_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|