clapclip.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "rich",
#     "torch",
#     "transformers",
#     "typer",
# ]
# ///

import os

os.environ["TOKENIZERS_PARALLELISM"] = "false"
import typer
from transformers import ClapModel, AutoTokenizer
from rich.prompt import Prompt
import json
import torch
import subprocess


def main(clap_model: str = "laion/clap-htsat-unfused"):
    model = ClapModel.from_pretrained(clap_model)
    tokenizer = AutoTokenizer.from_pretrained(clap_model)
    model.eval()
    while True:
        prompt = Prompt.ask("Prompt")
        if not prompt:
            break
        encoding = tokenizer([prompt], padding=True, return_tensors="pt")
        with torch.no_grad():
            output = model.get_text_features(**encoding).detach().numpy()[0]
        jemb = json.dumps(output.tolist())
        subprocess.run(["pbcopy"], input=jemb.encode("utf8"))
        print(f"embedded to clipboard")


if __name__ == "__main__":
    typer.run(main)