Add the export script

export.py

@@ -0,0 +1,140 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "coreai-core==1.0.0b3",
+#     "coreai-torch==0.4.3",
+#     "transformers==4.57.3",
+#     "torch<=2.13.0",
+# ]
+#
+# [tool.uv]
+# index-url       = "https://pypi.org/simple"
+# index-strategy  = "unsafe-best-match"
+# ///
+"""CLIP ViT-B/32 as one .aimodel with separate text and image encoders.
+
+Apple's models/clip recipe traces one function that needs an image and a fixed
+batch of 3 prompts on every call. Embedding text or images alone needs each
+tower on its own, so this exports:
+
+- encode_text(input_ids, attention_mask) -> text_embeds, inputs (B, 77) int32,
+  padded to CLIP's 77-token context
+- encode_image(pixel_values) -> image_embeds, input (B, 3, 224, 224)
+
+Both batch dims are dynamic (1-64), and the embeddings are L2-normalized.
+"""
+
+import argparse
+import shutil
+import time
+from pathlib import Path
+
+import torch
+import transformers
+from coreai.runtime import AIModelAssetMetadata
+from coreai_torch import TorchConverter, get_decomp_table
+
+MODEL = "openai/clip-vit-base-patch32"
+
+
+class TextEncoder(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.text_model = model.text_model
+        self.text_projection = model.text_projection
+
+    def forward(self, input_ids, attention_mask):
+        pooled = self.text_model(input_ids=input_ids, attention_mask=attention_mask)[1]
+        embeds = self.text_projection(pooled)
+        return embeds / embeds.norm(p=2, dim=-1, keepdim=True)
+
+
+class ImageEncoder(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.vision_model = model.vision_model
+        self.visual_projection = model.visual_projection
+
+    def forward(self, pixel_values):
+        pooled = self.vision_model(pixel_values=pixel_values)[1]
+        embeds = self.visual_projection(pooled)
+        return embeds / embeds.norm(p=2, dim=-1, keepdim=True)
+
+
+def export(module, kwargs, dynamic_shapes, dtype):
+    with torch.autocast(device_type="cpu", dtype=dtype):
+        exported = torch.export.export(module, args=(), kwargs=kwargs, dynamic_shapes=dynamic_shapes)
+    return exported.run_decompositions(get_decomp_table())
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--dtype", choices=["float16", "float32"], default="float16")
+    args = parser.parse_args()
+    dtype = getattr(torch, args.dtype)
+
+    model = transformers.CLIPModel.from_pretrained(MODEL).eval().to(dtype)
+    tokenizer = transformers.CLIPTokenizer.from_pretrained(MODEL)
+    text = tokenizer(
+        ["a photo of a cat", "a photo of a dog"],
+        return_tensors="pt",
+        padding="max_length",
+        max_length=77,
+    )
+    batch = torch.export.Dim("batch", min=1, max=64)
+
+    text_program = export(
+        TextEncoder(model),
+        {
+            "input_ids": text["input_ids"].to(torch.int32),
+            "attention_mask": text["attention_mask"].to(torch.int32),
+        },
+        {"input_ids": {0: batch}, "attention_mask": {0: batch}},
+        dtype,
+    )
+    image_program = export(
+        ImageEncoder(model),
+        {"pixel_values": torch.randn(2, 3, 224, 224).to(dtype)},
+        {"pixel_values": {0: batch}},
+        dtype,
+    )
+
+    converter = (
+        TorchConverter(mode=TorchConverter.Mode.RELEASE)
+        .add_exported_program(
+            exported_program=text_program,
+            input_names=["input_ids", "attention_mask"],
+            output_names=["text_embeds"],
+            entrypoint_name="encode_text",
+        )
+        .add_exported_program(
+            exported_program=image_program,
+            input_names=["pixel_values"],
+            output_names=["image_embeds"],
+            entrypoint_name="encode_image",
+        )
+    )
+    program = converter.to_coreai()
+
+    metadata = AIModelAssetMetadata()
+    metadata.author = "A. Radford et al."
+    metadata.license = "MIT"
+    metadata.model_description = (
+        "CLIP ViT-B/32 text and image encoders with L2-normalized embeddings. "
+        "Source: https://huggingface.co/openai/clip-vit-base-patch32"
+    )
+    metadata.creation_date = int(time.time())
+
+    path = Path(args.output_dir) / f"clip-vit-base-patch32-encoders_{args.dtype}.aimodel"
+    if path.exists():
+        shutil.rmtree(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    program.save_asset(path, metadata)
+    # The tokenizer travels with the model, so the app needs nothing else
+    tokenizer.save_pretrained(path.parent / "clip-tokenizer")
+    print(f"[INFO] Saved {path}")
+
+
+if __name__ == "__main__":
+    main()