Simple script to run the CoCa model

coca.py

from pathlib import Path
from sys import argv

import open_clip
import torch
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

coca_model, _, preprocess = open_clip.create_model_and_transforms(
    model_name="coca_ViT-L-14",
    pretrained="mscoco_finetuned_laion2B-s13B-b90k",
    device=device,
)


def generate_caption_coca(image):
    im = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad(), torch.autocast(device_type=device.type):
        generated = coca_model.generate(im, seq_len=20)
    return (
        open_clip.decode(generated[0].detach())
        .split("<end_of_text>")[0]
        .replace("<start_of_text>", "")
    )


assert len(argv) > 1, "No image path provided"
p = Path(argv[1]).absolute()
assert p.exists(), "No valid path provided"
im = Image.open(p).convert("RGB")
res = generate_caption_coca(im)
print(res)

CONDA_SETUP.md

conda create -n coca python
conda activate coca
conda install pytorch torchvision pytorch-cuda=11.7 -c pytorch -c nvidia
conda install transformers open-clip-torch -c conda-forge

CPU version:

conda create -n coca python
conda activate coca
conda install pytorch torchvision cpuonly -c pytorch
conda install transformers open-clip-torch -c conda-forge

• "a magnifying glass with a black handle . "

• "a cartoon of a pikachu wearing sunglasses ."

• "a few moments later title card"