# /// 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()