Clip

import { core as mx } from '@frost-beta/mlx';

export type ImageInputType = Buffer | ArrayBuffer | string;

export interface ProcessedImage {
    data: Buffer;
    info: sharp.OutputInfo;
}

export interface ClipInput {
    labels?: string[];
    images?: ProcessedImage[];
}

export interface ClipOutput {
    labelEmbeddings?: mx.array;
    imageEmbeddings?: mx.array;
}

export class Clip {
    constructor(modelDir: string);
    processImages(images: ImageInputType[]): Promise<ProcessedImage[]>;
    computeEmbeddings({ labels, images }: ClipInput): ClipOutput;
    /**
     * Short hands of computeEmbeddings to convert results to JavaScript numbers
     * and ensure the intermediate arrays are destroyed.
     */
    computeLabelEmbeddingsJs(labels: string[]): number[][];
    computeImageEmbeddingsJs(images: ProcessedImage[]): number[][];
    /**
     * Compute the cosine similarity between 2 embeddings.
     */
    static computeCosineSimilaritiy(a1: mx.array, a2: mx.array): mx.array;
    /**
     * Compute the cosine similarities between 2 arrays of embeddings.
     *
     * A tuple will be returned, with the first element being the cosine
     * similarity scores, and the second element being the indices sorted by
     * their scores from larger to smalller.
     */
    static computeCosineSimilarities(x1: mx.array | number[][], x2: mx.array | number[][]): [mx.array, mx.array];
}