| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """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) |
| |
| tokenizer.save_pretrained(path.parent / "clip-tokenizer") |
| print(f"[INFO] Saved {path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
