Add support for computing CLIP image and text embeddings separately (Closes huggingface#148) (huggingface#227)

* Define custom CLIP ONNX configs

* Update conversion script

* Support specifying custom model file name

* Use int64 for CLIP input ids

* Add support for CLIP text and vision models

* Fix JSDoc

* Add docs for `CLIPTextModelWithProjection`

* Add docs for `CLIPVisionModelWithProjection`

* Add unit test for CLIP text models

* Add unit test for CLIP vision models

* Set resize precision to 3 decimal places

* Fix `RawImage.save()` function

* Throw error when reading image and status != 200

* Create basic semantic image search application

* Separate out components

* Add `update-database` script

* Update transformers.js version

Author: xenova

examples/semantic-image-search/.env.local.example

@@ -0,0 +1,2 @@
+SUPABASE_URL=your-project-url
+SUPABASE_ANON_KEY=your-anon-key

examples/semantic-image-search/.eslintrc.json

@@ -0,0 +1,3 @@
+{
+  "extends": "next/core-web-vitals"
+}

examples/semantic-image-search/.gitignore

@@ -0,0 +1,35 @@
+# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
+
+# dependencies
+/node_modules
+/.pnp
+.pnp.js
+
+# testing
+/coverage
+
+# next.js
+/.next/
+/out/
+
+# production
+/build
+
+# misc
+.DS_Store
+*.pem
+
+# debug
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+
+# local env files
+.env*.local
+
+# vercel
+.vercel
+
+# typescript
+*.tsbuildinfo
+next-env.d.ts

examples/semantic-image-search/Dockerfile

@@ -0,0 +1,69 @@
+# syntax=docker/dockerfile:1.4
+
+# Adapted from https://github.com/vercel/next.js/blob/e60a1e747c3f521fc24dfd9ee2989e13afeb0a9b/examples/with-docker/Dockerfile
+# For more information, see https://nextjs.org/docs/pages/building-your-application/deploying#docker-image
+
+FROM node:18 AS base
+
+# Install dependencies only when needed
+FROM base AS deps
+WORKDIR /app
+
+# Install dependencies based on the preferred package manager
+COPY --link package.json yarn.lock* package-lock.json* pnpm-lock.yaml* ./
+RUN \
+  if [ -f yarn.lock ]; then yarn --frozen-lockfile; \
+  elif [ -f package-lock.json ]; then npm ci; \
+  elif [ -f pnpm-lock.yaml ]; then yarn global add pnpm && pnpm i --frozen-lockfile; \
+  else echo "Lockfile not found." && exit 1; \
+  fi
+
+
+# Rebuild the source code only when needed
+FROM base AS builder
+WORKDIR /app
+COPY --from=deps --link /app/node_modules ./node_modules
+COPY --link  . .
+
+# Next.js collects completely anonymous telemetry data about general usage.
+# Learn more here: https://nextjs.org/telemetry
+# Uncomment the following line in case you want to disable telemetry during the build.
+# ENV NEXT_TELEMETRY_DISABLED 1
+
+RUN npm run build
+
+# If using yarn comment out above and use below instead
+# RUN yarn build
+
+# Production image, copy all the files and run next
+FROM base AS runner
+WORKDIR /app
+
+ENV NODE_ENV production
+# Uncomment the following line in case you want to disable telemetry during runtime.
+# ENV NEXT_TELEMETRY_DISABLED 1
+
+RUN \
+  addgroup --system --gid 1001 nodejs; \
+  adduser --system --uid 1001 nextjs
+
+COPY --from=builder --link /app/public ./public
+
+# Automatically leverage output traces to reduce image size
+# https://nextjs.org/docs/advanced-features/output-file-tracing
+COPY --from=builder --link --chown=1001:1001 /app/.next/standalone ./
+COPY --from=builder --link --chown=1001:1001 /app/.next/static ./.next/static
+
+USER nextjs
+
+EXPOSE 3000
+
+ENV PORT 3000
+ENV HOSTNAME localhost
+
+# Allow the running process to write model files to the cache folder.
+# NOTE: In practice, you would probably want to pre-download the model files to avoid having to download them on-the-fly.
+RUN mkdir -p /app/node_modules/@xenova/.cache/
+RUN chmod 777 -R /app/node_modules/@xenova/
+
+CMD ["node", "server.js"]

examples/semantic-image-search/README.md

@@ -0,0 +1,34 @@
+This is a [Next.js](https://nextjs.org/) project bootstrapped with [`create-next-app`](https://github.com/vercel/next.js/tree/canary/packages/create-next-app).
+
+## Getting Started
+
+First, run the development server:
+
+```bash
+npm run dev
+# or
+yarn dev
+# or
+pnpm dev
+```
+
+Open [http://localhost:3000](http://localhost:3000) with your browser to see the result.
+
+You can start editing the page by modifying `app/page.js`. The page auto-updates as you edit the file.
+
+This project uses [`next/font`](https://nextjs.org/docs/basic-features/font-optimization) to automatically optimize and load Inter, a custom Google Font.
+
+## Learn More
+
+To learn more about Next.js, take a look at the following resources:
+
+- [Next.js Documentation](https://nextjs.org/docs) - learn about Next.js features and API.
+- [Learn Next.js](https://nextjs.org/learn) - an interactive Next.js tutorial.
+
+You can check out [the Next.js GitHub repository](https://github.com/vercel/next.js/) - your feedback and contributions are welcome!
+
+## Deploy on Vercel
+
+The easiest way to deploy your Next.js app is to use the [Vercel Platform](https://vercel.com/new?utm_medium=default-template&filter=next.js&utm_source=create-next-app&utm_campaign=create-next-app-readme) from the creators of Next.js.
+
+Check out our [Next.js deployment documentation](https://nextjs.org/docs/deployment) for more details.

examples/semantic-image-search/jsconfig.json

@@ -0,0 +1,7 @@
+{
+  "compilerOptions": {
+    "paths": {
+      "@/*": ["./src/*"]
+    }
+  }
+}

examples/semantic-image-search/next.config.js

@@ -0,0 +1,18 @@
+/** @type {import('next').NextConfig} */
+const nextConfig = {
+    // (Optional) Export as a standalone site
+    // See https://nextjs.org/docs/pages/api-reference/next-config-js/output#automatically-copying-traced-files
+    output: 'standalone', // Feel free to modify/remove this option
+
+    // Indicate that these packages should not be bundled by webpack
+    experimental: {
+        serverComponentsExternalPackages: ['sharp', 'onnxruntime-node'],
+    },
+
+    // Define which domains we are allowed to load images from
+    images: {
+        domains: ['images.unsplash.com'],
+    },
+};
+
+module.exports = nextConfig;

examples/semantic-image-search/package.json

@@ -0,0 +1,27 @@
+{
+  "name": "semantic-image-search",
+  "version": "0.1.0",
+  "private": true,
+  "scripts": {
+    "dev": "next dev",
+    "build": "next build",
+    "start": "next start",
+    "lint": "next lint"
+  },
+  "dependencies": {
+    "@xenova/transformers": "^2.5.0",
+    "@supabase/supabase-js": "^2.31.0",
+    "autoprefixer": "10.4.14",
+    "blurhash": "^2.0.5",
+    "eslint": "8.45.0",
+    "eslint-config-next": "13.4.12",
+    "next": "13.4.12",
+    "postcss": "8.4.27",
+    "react": "18.2.0",
+    "react-dom": "18.2.0",
+    "tailwindcss": "3.3.3"
+  },
+  "overrides": {
+    "protobufjs": "^7.2.4"
+  }
+}

examples/semantic-image-search/postcss.config.js

@@ -0,0 +1,6 @@
+module.exports = {
+  plugins: {
+    tailwindcss: {},
+    autoprefixer: {},
+  },
+}

examples/semantic-image-search/public/next.svg

@@ -0,0 +1,67 @@
+// Helper script to update the database with image embeddings
+
+import { AutoProcessor, RawImage, CLIPVisionModelWithProjection } from '@xenova/transformers';
+import { createClient } from '@supabase/supabase-js'
+
+if (!process.env.SUPABASE_SECRET_KEY) {
+    throw new Error('Missing `SUPABASE_SECRET_KEY` environment variable.')
+}
+
+// Create a single supabase client for interacting with your database
+const supabase = createClient(
+    process.env.SUPABASE_URL,
+    process.env.SUPABASE_SECRET_KEY,
+)
+
+let { data, error } = await supabase
+    .from('images')
+    .select('*')
+    .neq('ignore', true)
+    .is('image_embedding', null);
+
+if (error) {
+    throw error;
+}
+
+// Load processor and vision model
+const model_id = 'Xenova/clip-vit-base-patch16';
+const processor = await AutoProcessor.from_pretrained(model_id);
+const vision_model = await CLIPVisionModelWithProjection.from_pretrained(model_id, {
+    quantized: false,
+});
+
+for (const image_data of data) {
+    let image;
+    try {
+        image = await RawImage.read(image_data.photo_image_url);
+    } catch (e) {
+        // Unable to load image, so we ignore it
+        console.warn('Ignoring image due to error', e)
+        await supabase
+            .from('images')
+            .update({ ignore: true })
+            .eq('photo_id', image_data.photo_id)
+            .select()
+        continue;
+    }
+
+    // Read image and run processor
+    let image_inputs = await processor(image);
+
+    // Compute embeddings
+    const { image_embeds } = await vision_model(image_inputs);
+    const embed_as_list = image_embeds.tolist()[0];
+
+    // https://supabase.com/docs/guides/ai/vector-columns#storing-a-vector--embedding
+    const { data, error } = await supabase
+        .from('images')
+        .update({ image_embedding: embed_as_list })
+        .eq('photo_id', image_data.photo_id)
+        .select()
+
+    if (error) {
+        console.error('error', error)
+    } else {
+        console.log('success', image_data.photo_id)
+    }
+}

examples/semantic-image-search/public/vercel.svg

@@ -0,0 +1,51 @@
+import { AutoTokenizer, CLIPTextModelWithProjection } from "@xenova/transformers";
+import { createClient } from '@supabase/supabase-js'
+
+// Use the Singleton pattern to enable lazy construction of the pipeline.
+// NOTE: We wrap the class in a function to prevent code duplication (see below).
+const S = () => class ApplicationSingleton {
+    static model_id = 'Xenova/clip-vit-base-patch16';
+    static tokenizer = null;
+    static text_model = null;
+    static database = null;
+
+    static async getInstance() {
+        // Load tokenizer and text model
+        if (this.tokenizer === null) {
+            this.tokenizer = AutoTokenizer.from_pretrained(this.model_id);
+        }
+
+        if (this.text_model === null) {
+            this.text_model = CLIPTextModelWithProjection.from_pretrained(this.model_id, {
+                quantized: false,
+            });
+        }
+
+        if (this.database === null) {
+            this.database = createClient(
+                process.env.SUPABASE_URL,
+                process.env.SUPABASE_ANON_KEY,
+            )
+        }
+
+        return Promise.all([
+            this.tokenizer,
+            this.text_model,
+            this.database,
+        ]);
+    }
+}
+
+let ApplicationSingleton;
+if (process.env.NODE_ENV !== 'production') {
+    // When running in development mode, attach the pipeline to the
+    // global object so that it's preserved between hot reloads.
+    // For more information, see https://vercel.com/guides/nextjs-prisma-postgres
+    if (!global.ApplicationSingleton) {
+        global.ApplicationSingleton = S();
+    }
+    ApplicationSingleton = global.ApplicationSingleton;
+} else {
+    ApplicationSingleton = S();
+}
+export default ApplicationSingleton;

examples/semantic-image-search/scripts/update-database.mjs

@@ -0,0 +1,52 @@
+import Image from 'next/image'
+import { blurHashToDataURL } from '../utils.js'
+
+export function ImageGrid({ images, setCurrentImage }) {
+    return (
+        <div className="columns-2 gap-4 sm:columns-3 xl:columns-4 2xl:columns-5">
+            {images && images.map(({
+                photo_id,
+                photo_url,
+                photo_image_url,
+                photo_aspect_ratio,
+                photo_width,
+                photo_height,
+                blur_hash,
+                photo_description,
+                ai_description,
+                similarity,
+            }) => (
+                <div
+                    key={photo_id}
+                    href={photo_url}
+                    className='after:content group cursor-pointer relative mb-4 block w-full after:pointer-events-none after:absolute after:inset-0 after:rounded-lg after:shadow-highlight'
+                    onClick={() => {
+                        setCurrentImage({
+                            photo_id,
+                            photo_url,
+                            photo_image_url,
+                            photo_aspect_ratio,
+                            photo_width,
+                            photo_height,
+                            blur_hash,
+                            photo_description,
+                            ai_description,
+                            similarity,
+                        });
+                    }}
+                >
+                    <Image
+                        alt={photo_description || ai_description || ""}
+                        className="transform rounded-lg brightness-90 transition will-change-auto group-hover:brightness-110"
+                        style={{ transform: 'translate3d(0, 0, 0)' }}
+                        placeholder="blur"
+                        blurDataURL={blurHashToDataURL(blur_hash)}
+                        src={`${photo_image_url}?auto=format&fit=crop&w=480&q=80`}
+                        width={480}
+                        height={480 / photo_aspect_ratio}
+                        unoptimized={true}
+                    />
+                </div>
+            ))}
+        </div>)
+}