Update cog

Recommending the stable 0.5.4 version for now. 
0.6.0 supports BLIP2 but requires newer transformers version and could use more love still.

.github/workflows/python-publish.yml

@@ -0,0 +1,39 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: pharmapsychotic
+        password: ${{ secrets.PYPI_API_TOKEN }}

.gitignore

@@ -3,6 +3,7 @@
 .vscode/
 bench/
 cache/
+ci_env/
 clip-interrogator/
 clip_interrogator.egg-info/
 dist/

MANIFEST.in

@@ -2,4 +2,5 @@ include clip_interrogator/data/artists.txt
 include clip_interrogator/data/flavors.txt
 include clip_interrogator/data/mediums.txt
 include clip_interrogator/data/movements.txt
+include clip_interrogator/data/negative.txt
 include requirements.txt

README.md

@@ -4,18 +4,17 @@
 
 ## Run it!
 
-Run Version 2 on Colab, HuggingFace, and Replicate!
-
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/main/clip_interrogator.ipynb) [![Generic badge](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue.svg)](https://huggingface.co/spaces/pharma/CLIP-Interrogator) [![Replicate](https://replicate.com/cjwbw/clip-interrogator/badge)](https://replicate.com/cjwbw/clip-interrogator)
+🆕 Now available as a [Stable Diffusion Web UI Extension](https://github.com/pharmapsychotic/clip-interrogator-ext)! 🆕
 
 <br>
 
-For **Stable Diffusion 2.0** prompting use the `ViT-H` version:
+Run Version 2 on Colab, HuggingFace, and Replicate!
 
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/open-clip/clip_interrogator.ipynb) [![Generic badge](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue.svg)](https://huggingface.co/spaces/fffiloni/CLIP-Interrogator-2)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/main/clip_interrogator.ipynb) [![Generic badge](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue.svg)](https://huggingface.co/spaces/pharma/CLIP-Interrogator) [![Replicate](https://replicate.com/pharmapsychotic/clip-interrogator/badge)](https://replicate.com/pharmapsychotic/clip-interrogator) [![Lambda](https://img.shields.io/badge/%CE%BB-Lambda-blue)](https://cloud.lambdalabs.com/demos/ml/CLIP-Interrogator)
 
 <br>
 
+
 Version 1 still available in Colab for comparing different CLIP models 
 
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/v1/clip_interrogator.ipynb) 
@@ -31,19 +30,26 @@ The **CLIP Interrogator** is a prompt engineering tool that combines OpenAI's [C
 Create and activate a Python virtual environment
 ```bash
 python3 -m venv ci_env
-source ci_env/bin/activate
+(for linux  ) source ci_env/bin/activate
+(for windows) .\ci_env\Scripts\activate
 ```
 
 Install with PIP
 ```
-pip install -e git+https://github.com/pharmapsychotic/BLIP.git@lib#egg=blip
-pip install clip-interrogator
+# install torch with GPU support for example:
+pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu117
+
+# install clip-interrogator
+pip install clip-interrogator==0.5.4
+
+# or for very latest WIP with BLIP2 support
+#pip install clip-interrogator==0.6.0
 ```
 
 You can then use it in your script
 ```python
 from PIL import Image
-from clip_interrogator import Interrogator, Config
+from clip_interrogator import Config, Interrogator
 image = Image.open(image_path).convert('RGB')
 ci = Interrogator(Config(clip_model_name="ViT-L-14/openai"))
 print(ci.interrogate(image))
@@ -52,3 +58,29 @@ print(ci.interrogate(image))
 CLIP Interrogator uses OpenCLIP which supports many different pretrained CLIP models. For the best prompts for 
 Stable Diffusion 1.X use `ViT-L-14/openai` for clip_model_name. For Stable Diffusion 2.0 use `ViT-H-14/laion2b_s32b_b79k`
 
+## Configuration
+
+The `Config` object lets you configure CLIP Interrogator's processing. 
+* `clip_model_name`: which of the OpenCLIP pretrained CLIP models to use
+* `cache_path`: path where to save precomputed text embeddings
+* `download_cache`: when True will download the precomputed embeddings from huggingface
+* `chunk_size`: batch size for CLIP, use smaller for lower VRAM
+* `quiet`: when True no progress bars or text output will be displayed
+
+On systems with low VRAM you can call `config.apply_low_vram_defaults()` to reduce the amount of VRAM needed (at the cost of some speed and quality). The default settings use about 6.3GB of VRAM and the low VRAM settings use about 2.7GB.
+
+See the [run_cli.py](https://github.com/pharmapsychotic/clip-interrogator/blob/main/run_cli.py) and [run_gradio.py](https://github.com/pharmapsychotic/clip-interrogator/blob/main/run_gradio.py) for more examples on using Config and Interrogator classes.
+
+
+## Ranking against your own list of terms (requires version 0.6.0)
+
+```python
+from clip_interrogator import Config, Interrogator, LabelTable, load_list
+from PIL import Image
+
+ci = Interrogator(Config(blip_model_type=None))
+image = Image.open(image_path).convert('RGB')
+table = LabelTable(load_list('terms.txt'), 'terms', ci)
+best_match = table.rank(ci.image_to_features(image), top_count=1)[0]
+print(best_match)
+```

clip_interrogator.ipynb

@@ -1,21 +1,35 @@
 {
   "cells": [
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3jm8RYrLqvzz"
       },
       "source": [
-        "# CLIP Interrogator 2.1 ViT-H special edition!\n",
+        "# CLIP Interrogator 2.4 by [@pharmapsychotic](https://twitter.com/pharmapsychotic) \n",
         "\n",
         "Want to figure out what a good prompt might be to create new images like an existing one? The CLIP Interrogator is here to get you answers!\n",
         "\n",
-        "This version is specialized for producing nice prompts for use with **[Stable Diffusion 2.0](https://stability.ai/blog/stable-diffusion-v2-release)** using the **ViT-H-14** OpenCLIP model!\n"
+        "<br>\n",
+        "\n",
+        "For Stable Diffusion 1.X choose the **ViT-L** model and for Stable Diffusion 2.0+ choose the **ViT-H** CLIP Model.\n",
+        "\n",
+        "This version is specialized for producing nice prompts for use with Stable Diffusion and achieves higher alignment between generated text prompt and source image. You can try out the old [version 1](https://colab.research.google.com/github/pharmapsychotic/clip-interrogator/blob/v1/clip_interrogator.ipynb) to see how different CLIP models ranks terms. \n",
+        "\n",
+        "You can also run this on HuggingFace and Replicate<br>\n",
+        "[![Generic badge](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue.svg)](https://huggingface.co/spaces/pharma/CLIP-Interrogator) [![Replicate](https://replicate.com/pharmapsychotic/clip-interrogator/badge)](https://replicate.com/pharmapsychotic/clip-interrogator)\n",
+        "\n",
+        "<br>\n",
+        "\n",
+        "If this notebook is helpful to you please consider buying me a coffee via [ko-fi](https://ko-fi.com/pharmapsychotic) or following me on [twitter](https://twitter.com/pharmapsychotic) for more cool Ai stuff. 🙂\n",
+        "\n",
+        "And if you're looking for more Ai art tools check out my [Ai generative art tools list](https://pharmapsychotic.com/tools.html).\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
       "metadata": {
         "cellView": "form",
         "id": "aP9FjmWxtLKJ"
@@ -28,7 +42,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": null,
       "metadata": {
         "cellView": "form",
         "id": "xpPKQR40qvz2"
@@ -40,57 +54,63 @@
         "\n",
         "def setup():\n",
         "    install_cmds = [\n",
-        "        ['pip', 'install', 'ftfy', 'gradio', 'regex', 'tqdm', 'transformers==4.21.2', 'timm', 'fairscale', 'requests'],\n",
+        "        ['pip', 'install', 'gradio'],\n",
         "        ['pip', 'install', 'open_clip_torch'],\n",
-        "        ['pip', 'install', '-e', 'git+https://github.com/pharmapsychotic/BLIP.git@lib#egg=blip'],\n",
-        "        ['git', 'clone', '-b', 'open-clip', 'https://github.com/pharmapsychotic/clip-interrogator.git']\n",
+        "        ['pip', 'install', 'clip-interrogator'],\n",
         "    ]\n",
         "    for cmd in install_cmds:\n",
         "        print(subprocess.run(cmd, stdout=subprocess.PIPE).stdout.decode('utf-8'))\n",
         "\n",
         "setup()\n",
         "\n",
-        "# download cache files\n",
-        "print(\"Download preprocessed cache files...\")\n",
-        "CACHE_URLS = [\n",
-        "    'https://huggingface.co/pharma/ci-preprocess/resolve/main/ViT-H-14_laion2b_s32b_b79k_artists.pkl',\n",
-        "    'https://huggingface.co/pharma/ci-preprocess/resolve/main/ViT-H-14_laion2b_s32b_b79k_flavors.pkl',\n",
-        "    'https://huggingface.co/pharma/ci-preprocess/resolve/main/ViT-H-14_laion2b_s32b_b79k_mediums.pkl',\n",
-        "    'https://huggingface.co/pharma/ci-preprocess/resolve/main/ViT-H-14_laion2b_s32b_b79k_movements.pkl',\n",
-        "    'https://huggingface.co/pharma/ci-preprocess/resolve/main/ViT-H-14_laion2b_s32b_b79k_trendings.pkl',\n",
-        "]\n",
-        "os.makedirs('cache', exist_ok=True)\n",
-        "for url in CACHE_URLS:\n",
-        "    print(subprocess.run(['wget', url, '-P', 'cache'], stdout=subprocess.PIPE).stdout.decode('utf-8'))\n",
-        "\n",
-        "import sys\n",
-        "sys.path.append('src/blip')\n",
-        "sys.path.append('clip-interrogator')\n",
+        "\n",
+        "caption_model_name = 'blip-large' #@param [\"blip-base\", \"blip-large\", \"git-large-coco\"]\n",
+        "clip_model_name = 'ViT-L-14/openai' #@param [\"ViT-L-14/openai\", \"ViT-H-14/laion2b_s32b_b79k\"]\n",
         "\n",
         "import gradio as gr\n",
         "from clip_interrogator import Config, Interrogator\n",
         "\n",
         "config = Config()\n",
-        "config.blip_offload = True\n",
-        "config.chunk_size = 2048\n",
-        "config.flavor_intermediate_count = 512\n",
-        "config.blip_num_beams = 64\n",
-        "\n",
+        "config.clip_model_name = clip_model_name\n",
+        "config.caption_model_name = caption_model_name\n",
         "ci = Interrogator(config)\n",
         "\n",
-        "def inference(image, mode, best_max_flavors):\n",
+        "def image_analysis(image):\n",
+        "    image = image.convert('RGB')\n",
+        "    image_features = ci.image_to_features(image)\n",
+        "\n",
+        "    top_mediums = ci.mediums.rank(image_features, 5)\n",
+        "    top_artists = ci.artists.rank(image_features, 5)\n",
+        "    top_movements = ci.movements.rank(image_features, 5)\n",
+        "    top_trendings = ci.trendings.rank(image_features, 5)\n",
+        "    top_flavors = ci.flavors.rank(image_features, 5)\n",
+        "\n",
+        "    medium_ranks = {medium: sim for medium, sim in zip(top_mediums, ci.similarities(image_features, top_mediums))}\n",
+        "    artist_ranks = {artist: sim for artist, sim in zip(top_artists, ci.similarities(image_features, top_artists))}\n",
+        "    movement_ranks = {movement: sim for movement, sim in zip(top_movements, ci.similarities(image_features, top_movements))}\n",
+        "    trending_ranks = {trending: sim for trending, sim in zip(top_trendings, ci.similarities(image_features, top_trendings))}\n",
+        "    flavor_ranks = {flavor: sim for flavor, sim in zip(top_flavors, ci.similarities(image_features, top_flavors))}\n",
+        "    \n",
+        "    return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks\n",
+        "\n",
+        "def image_to_prompt(image, mode):\n",
+        "    ci.config.chunk_size = 2048 if ci.config.clip_model_name == \"ViT-L-14/openai\" else 1024\n",
+        "    ci.config.flavor_intermediate_count = 2048 if ci.config.clip_model_name == \"ViT-L-14/openai\" else 1024\n",
         "    image = image.convert('RGB')\n",
         "    if mode == 'best':\n",
-        "        return ci.interrogate(image, max_flavors=int(best_max_flavors))\n",
+        "        return ci.interrogate(image)\n",
         "    elif mode == 'classic':\n",
         "        return ci.interrogate_classic(image)\n",
-        "    else:\n",
-        "        return ci.interrogate_fast(image)\n"
+        "    elif mode == 'fast':\n",
+        "        return ci.interrogate_fast(image)\n",
+        "    elif mode == 'negative':\n",
+        "        return ci.interrogate_negative(image)\n",
+        "        "
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": null,
       "metadata": {
         "cellView": "form",
         "colab": {
@@ -100,59 +120,40 @@
         "id": "Pf6qkFG6MPRj",
         "outputId": "8d542b56-8be7-453d-bf27-d0540a774c7d"
       },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Colab notebook detected. To show errors in colab notebook, set `debug=True` in `launch()`\n",
-            "\n",
-            "Using Embedded Colab Mode (NEW). If you have issues, please use share=True and file an issue at https://github.com/gradio-app/gradio/\n",
-            "Note: opening the browser inspector may crash Embedded Colab Mode.\n",
-            "\n",
-            "To create a public link, set `share=True` in `launch()`.\n"
-          ]
-        },
-        {
-          "data": {
-            "application/javascript": "(async (port, path, width, height, cache, element) => {\n                        if (!google.colab.kernel.accessAllowed && !cache) {\n                            return;\n                        }\n                        element.appendChild(document.createTextNode(''));\n                        const url = await google.colab.kernel.proxyPort(port, {cache});\n\n                        const external_link = document.createElement('div');\n                        external_link.innerHTML = `\n                            <div style=\"font-family: monospace; margin-bottom: 0.5rem\">\n                                Running on <a href=${new URL(path, url).toString()} target=\"_blank\">\n                                    https://localhost:${port}${path}\n                                </a>\n                            </div>\n                        `;\n                        element.appendChild(external_link);\n\n                        const iframe = document.createElement('iframe');\n                        iframe.src = new URL(path, url).toString();\n                        iframe.height = height;\n                        iframe.allow = \"autoplay; camera; microphone; clipboard-read; clipboard-write;\"\n                        iframe.width = width;\n                        iframe.style.border = 0;\n                        element.appendChild(iframe);\n                    })(7860, \"/\", \"100%\", 500, false, window.element)",
-            "text/plain": [
-              "<IPython.core.display.Javascript object>"
-            ]
-          },
-          "metadata": {},
-          "output_type": "display_data"
-        },
-        {
-          "data": {
-            "text/plain": [
-              "(<gradio.routes.App at 0x7f894e553710>, 'http://127.0.0.1:7860/', None)"
-            ]
-          },
-          "execution_count": 4,
-          "metadata": {},
-          "output_type": "execute_result"
-        }
-      ],
+      "outputs": [],
       "source": [
         "#@title Image to prompt! 🖼️ -> 📝\n",
         "   \n",
-        "inputs = [\n",
-        "    gr.inputs.Image(type='pil'),\n",
-        "    gr.Radio(['best', 'classic', 'fast'], label='', value='best'),\n",
-        "    gr.Number(value=4, label='best mode max flavors'),\n",
-        "]\n",
-        "outputs = [\n",
-        "    gr.outputs.Textbox(label=\"Output\"),\n",
-        "]\n",
-        "\n",
-        "io = gr.Interface(\n",
-        "    inference, \n",
-        "    inputs, \n",
-        "    outputs, \n",
-        "    allow_flagging=False,\n",
-        ")\n",
-        "io.launch()\n"
+        "def prompt_tab():\n",
+        "    with gr.Column():\n",
+        "        with gr.Row():\n",
+        "            image = gr.Image(type='pil', label=\"Image\")\n",
+        "            with gr.Column():\n",
+        "                mode = gr.Radio(['best', 'fast', 'classic', 'negative'], label='Mode', value='best')\n",
+        "        prompt = gr.Textbox(label=\"Prompt\")\n",
+        "    button = gr.Button(\"Generate prompt\")\n",
+        "    button.click(image_to_prompt, inputs=[image, mode], outputs=prompt)\n",
+        "\n",
+        "def analyze_tab():\n",
+        "    with gr.Column():\n",
+        "        with gr.Row():\n",
+        "            image = gr.Image(type='pil', label=\"Image\")\n",
+        "        with gr.Row():\n",
+        "            medium = gr.Label(label=\"Medium\", num_top_classes=5)\n",
+        "            artist = gr.Label(label=\"Artist\", num_top_classes=5)        \n",
+        "            movement = gr.Label(label=\"Movement\", num_top_classes=5)\n",
+        "            trending = gr.Label(label=\"Trending\", num_top_classes=5)\n",
+        "            flavor = gr.Label(label=\"Flavor\", num_top_classes=5)\n",
+        "    button = gr.Button(\"Analyze\")\n",
+        "    button.click(image_analysis, inputs=image, outputs=[medium, artist, movement, trending, flavor])\n",
+        "\n",
+        "with gr.Blocks() as ui:\n",
+        "    with gr.Tab(\"Prompt\"):\n",
+        "        prompt_tab()\n",
+        "    with gr.Tab(\"Analyze\"):\n",
+        "        analyze_tab()\n",
+        "\n",
+        "ui.launch(show_api=False, debug=False)\n"
       ]
     },
     {
@@ -179,10 +180,9 @@
         "from tqdm import tqdm\n",
         "\n",
         "folder_path = \"/content/my_images\" #@param {type:\"string\"}\n",
-        "prompt_mode = 'best' #@param [\"best\",\"classic\", \"fast\"]\n",
+        "prompt_mode = 'best' #@param [\"best\",\"fast\",\"classic\",\"negative\"]\n",
         "output_mode = 'rename' #@param [\"desc.csv\",\"rename\"]\n",
         "max_filename_len = 128 #@param {type:\"integer\"}\n",
-        "best_max_flavors = 4 #@param {type:\"integer\"}\n",
         "\n",
         "\n",
         "def sanitize_for_filename(prompt: str, max_len: int) -> str:\n",
@@ -199,7 +199,7 @@
         "        clear_output(wait=True)\n",
         "\n",
         "    image = Image.open(os.path.join(folder_path, file)).convert('RGB')\n",
-        "    prompt = inference(image, prompt_mode, best_max_flavors=best_max_flavors)\n",
+        "    prompt = image_to_prompt(image, prompt_mode)\n",
         "    prompts.append(prompt)\n",
         "\n",
         "    print(prompt)\n",
@@ -242,7 +242,7 @@
       "provenance": []
     },
     "kernelspec": {
-      "display_name": "Python 3.8.10 ('venv': venv)",
+      "display_name": "Python 3.7.15 ('py37')",
       "language": "python",
       "name": "python3"
     },
@@ -256,12 +256,12 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.8.10"
+      "version": "3.9.5"
     },
     "orig_nbformat": 4,
     "vscode": {
       "interpreter": {
-        "hash": "f7a8d9541664ade9cff251487a19c76f2dd1b4c864d158f07ee26d1b0fd5c9a1"
+        "hash": "1f51d5616d3bc2b87a82685314c5be1ec9a49b6e0cb1f707bfa2acb6c45f3e5f"
       }
     }
   },

clip_interrogator/__init__.py

@@ -1,4 +1,4 @@
-from .clip_interrogator import Interrogator, Config
+from .clip_interrogator import Config, Interrogator, LabelTable, list_caption_models, list_clip_models, load_list
 
-__version__ = '0.2.0'
+__version__ = '0.6.0'
 __author__ = 'pharmapsychotic'

clip_interrogator/clip_interrogator.py

@@ -1,148 +1,211 @@
 import hashlib
-import inspect
 import math
 import numpy as np
 import open_clip
 import os
-import pickle
+import requests
 import time
 import torch
 
 from dataclasses import dataclass
-from models.blip import blip_decoder
 from PIL import Image
-from torchvision import transforms
-from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, Blip2ForConditionalGeneration
 from tqdm import tqdm
-from typing import List
+from typing import List, Optional
+
+from safetensors.numpy import load_file, save_file
+
+CAPTION_MODELS = {
+    'blip-base': 'Salesforce/blip-image-captioning-base',   # 990MB
+    'blip-large': 'Salesforce/blip-image-captioning-large', # 1.9GB
+    'blip2-2.7b': 'Salesforce/blip2-opt-2.7b',              # 15.5GB
+    'blip2-flan-t5-xl': 'Salesforce/blip2-flan-t5-xl',      # 15.77GB
+    'git-large-coco': 'microsoft/git-large-coco',           # 1.58GB
+}
+
+CACHE_URL_BASE = 'https://huggingface.co/pharmapsychotic/ci-preprocess/resolve/main/'
 
 
 @dataclass 
 class Config:
     # models can optionally be passed in directly
-    blip_model = None
+    caption_model = None
+    caption_processor = None
     clip_model = None
     clip_preprocess = None
 
     # blip settings
-    blip_image_eval_size: int = 384
-    blip_max_length: int = 32
-    blip_model_url: str = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth'
-    blip_num_beams: int = 8
-    blip_offload: bool = False
+    caption_max_length: int = 32
+    caption_model_name: Optional[str] = 'blip-large' # use a key from CAPTION_MODELS or None
+    caption_offload: bool = False
 
     # clip settings
-    clip_model_name: str = 'ViT-H-14/laion2b_s32b_b79k'
-    clip_model_path: str = None
+    clip_model_name: str = 'ViT-L-14/openai'
+    clip_model_path: Optional[str] = None
+    clip_offload: bool = False
 
     # interrogator settings
-    cache_path: str = 'cache'
-    chunk_size: int = 2048
+    cache_path: str = 'cache'   # path to store cached text embeddings
+    download_cache: bool = True # when true, cached embeds are downloaded from huggingface
+    chunk_size: int = 2048      # batch size for CLIP, use smaller for lower VRAM
     data_path: str = os.path.join(os.path.dirname(__file__), 'data')
-    device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device: str = ("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
     flavor_intermediate_count: int = 2048
     quiet: bool = False # when quiet progress bars are not shown
 
+    def apply_low_vram_defaults(self):
+        self.caption_model_name = 'blip-base'
+        self.caption_offload = True
+        self.clip_offload = True
+        self.chunk_size = 1024
+        self.flavor_intermediate_count = 1024
 
 class Interrogator():
     def __init__(self, config: Config):
         self.config = config
         self.device = config.device
+        self.dtype = torch.float16 if self.device == 'cuda' else torch.float32
+        self.caption_offloaded = True
+        self.clip_offloaded = True
+        self.load_caption_model()
+        self.load_clip_model()
 
-        if config.blip_model is None:
-            if not config.quiet:
-                print("Loading BLIP model...")
-            blip_path = os.path.dirname(inspect.getfile(blip_decoder))
-            configs_path = os.path.join(os.path.dirname(blip_path), 'configs')
-            med_config = os.path.join(configs_path, 'med_config.json')
-            blip_model = blip_decoder(
-                pretrained=config.blip_model_url, 
-                image_size=config.blip_image_eval_size, 
-                vit='large', 
-                med_config=med_config
-            )
-            blip_model.eval()
-            blip_model = blip_model.to(config.device)
-            self.blip_model = blip_model
+    def load_caption_model(self):
+        if self.config.caption_model is None and self.config.caption_model_name:
+            if not self.config.quiet:
+                print(f"Loading caption model {self.config.caption_model_name}...")
+
+            model_path = CAPTION_MODELS[self.config.caption_model_name]
+            if self.config.caption_model_name.startswith('git-'):
+                caption_model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float32)
+            elif self.config.caption_model_name.startswith('blip2-'):
+                caption_model = Blip2ForConditionalGeneration.from_pretrained(model_path, torch_dtype=self.dtype)
+            else:
+                caption_model = BlipForConditionalGeneration.from_pretrained(model_path, torch_dtype=self.dtype)
+            self.caption_processor = AutoProcessor.from_pretrained(model_path)
+
+            caption_model.eval()
+            if not self.config.caption_offload:
+                caption_model = caption_model.to(self.config.device)
+            self.caption_model = caption_model
         else:
-            self.blip_model = config.blip_model
-
-        self.load_clip_model()
+            self.caption_model = self.config.caption_model
+            self.caption_processor = self.config.caption_processor
 
     def load_clip_model(self):
         start_time = time.time()
         config = self.config
 
+        clip_model_name, clip_model_pretrained_name = config.clip_model_name.split('/', 2)
+
         if config.clip_model is None:
             if not config.quiet:
-                print("Loading CLIP model...")
+                print(f"Loading CLIP model {config.clip_model_name}...")
 
-            clip_model_name, clip_model_pretrained_name = config.clip_model_name.split('/', 2)
             self.clip_model, _, self.clip_preprocess = open_clip.create_model_and_transforms(
                 clip_model_name, 
                 pretrained=clip_model_pretrained_name, 
-                precision='fp16',
+                precision='fp16' if config.device == 'cuda' else 'fp32',
                 device=config.device,
                 jit=False,
                 cache_dir=config.clip_model_path
             )
-            self.clip_model.half().to(config.device).eval()
+            self.clip_model.eval()
         else:
             self.clip_model = config.clip_model
             self.clip_preprocess = config.clip_preprocess
         self.tokenize = open_clip.get_tokenizer(clip_model_name)
 
-        sites = ['Artstation', 'behance', 'cg society', 'cgsociety', 'deviantart', 'dribble', 'flickr', 'instagram', 'pexels', 'pinterest', 'pixabay', 'pixiv', 'polycount', 'reddit', 'shutterstock', 'tumblr', 'unsplash', 'zbrush central']
+        sites = ['Artstation', 'behance', 'cg society', 'cgsociety', 'deviantart', 'dribbble', 
+                 'flickr', 'instagram', 'pexels', 'pinterest', 'pixabay', 'pixiv', 'polycount', 
+                 'reddit', 'shutterstock', 'tumblr', 'unsplash', 'zbrush central']
         trending_list = [site for site in sites]
         trending_list.extend(["trending on "+site for site in sites])
         trending_list.extend(["featured on "+site for site in sites])
         trending_list.extend([site+" contest winner" for site in sites])
 
-        raw_artists = _load_list(config.data_path, 'artists.txt')
+        raw_artists = load_list(config.data_path, 'artists.txt')
         artists = [f"by {a}" for a in raw_artists]
         artists.extend([f"inspired by {a}" for a in raw_artists])
 
-        self.artists = LabelTable(artists, "artists", self.clip_model, self.tokenize, config)
-        self.flavors = LabelTable(_load_list(config.data_path, 'flavors.txt'), "flavors", self.clip_model, self.tokenize, config)
-        self.mediums = LabelTable(_load_list(config.data_path, 'mediums.txt'), "mediums", self.clip_model, self.tokenize, config)
-        self.movements = LabelTable(_load_list(config.data_path, 'movements.txt'), "movements", self.clip_model, self.tokenize, config)
-        self.trendings = LabelTable(trending_list, "trendings", self.clip_model, self.tokenize, config)
+        self._prepare_clip()
+        self.artists = LabelTable(artists, "artists", self)
+        self.flavors = LabelTable(load_list(config.data_path, 'flavors.txt'), "flavors", self)
+        self.mediums = LabelTable(load_list(config.data_path, 'mediums.txt'), "mediums", self)
+        self.movements = LabelTable(load_list(config.data_path, 'movements.txt'), "movements", self)
+        self.trendings = LabelTable(trending_list, "trendings", self)
+        self.negative = LabelTable(load_list(config.data_path, 'negative.txt'), "negative", self)
 
         end_time = time.time()
         if not config.quiet:
             print(f"Loaded CLIP model and data in {end_time-start_time:.2f} seconds.")
 
+    def chain(
+        self, 
+        image_features: torch.Tensor, 
+        phrases: List[str], 
+        best_prompt: str="", 
+        best_sim: float=0, 
+        min_count: int=8,
+        max_count: int=32, 
+        desc="Chaining", 
+        reverse: bool=False
+    ) -> str:
+        self._prepare_clip()
+
+        phrases = set(phrases)
+        if not best_prompt:
+            best_prompt = self.rank_top(image_features, [f for f in phrases], reverse=reverse)
+            best_sim = self.similarity(image_features, best_prompt)
+            phrases.remove(best_prompt)
+        curr_prompt, curr_sim = best_prompt, best_sim
+        
+        def check(addition: str, idx: int) -> bool:
+            nonlocal best_prompt, best_sim, curr_prompt, curr_sim
+            prompt = curr_prompt + ", " + addition
+            sim = self.similarity(image_features, prompt)
+            if reverse:
+                sim = -sim
+            
+            if sim > best_sim:
+                best_prompt, best_sim = prompt, sim
+            if sim > curr_sim or idx < min_count:
+                curr_prompt, curr_sim = prompt, sim
+                return True
+            return False
+
+        for idx in tqdm(range(max_count), desc=desc, disable=self.config.quiet):
+            best = self.rank_top(image_features, [f"{curr_prompt}, {f}" for f in phrases], reverse=reverse)
+            flave = best[len(curr_prompt)+2:]
+            if not check(flave, idx):
+                break
+            if _prompt_at_max_len(curr_prompt, self.tokenize):
+                break
+            phrases.remove(flave)
+
+        return best_prompt
+
     def generate_caption(self, pil_image: Image) -> str:
-        if self.config.blip_offload:
-            self.blip_model = self.blip_model.to(self.device)
-        size = self.config.blip_image_eval_size
-        gpu_image = transforms.Compose([
-            transforms.Resize((size, size), interpolation=InterpolationMode.BICUBIC),
-            transforms.ToTensor(),
-            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
-        ])(pil_image).unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            caption = self.blip_model.generate(
-                gpu_image, 
-                sample=False, 
-                num_beams=self.config.blip_num_beams, 
-                max_length=self.config.blip_max_length, 
-                min_length=5
-            )
-        if self.config.blip_offload:
-            self.blip_model = self.blip_model.to("cpu")
-        return caption[0]
+        assert self.caption_model is not None, "No caption model loaded."
+        self._prepare_caption()
+        inputs = self.caption_processor(images=pil_image, return_tensors="pt").to(self.device)
+        if not self.config.caption_model_name.startswith('git-'):
+            inputs = inputs.to(self.dtype)
+        tokens = self.caption_model.generate(**inputs, max_new_tokens=self.config.caption_max_length)
+        return self.caption_processor.batch_decode(tokens, skip_special_tokens=True)[0].strip()
 
     def image_to_features(self, image: Image) -> torch.Tensor:
+        self._prepare_clip()
         images = self.clip_preprocess(image).unsqueeze(0).to(self.device)
         with torch.no_grad(), torch.cuda.amp.autocast():
             image_features = self.clip_model.encode_image(images)
             image_features /= image_features.norm(dim=-1, keepdim=True)
         return image_features
 
-    def interrogate_classic(self, image: Image, max_flavors: int=3) -> str:
-        caption = self.generate_caption(image)
+    def interrogate_classic(self, image: Image, max_flavors: int=3, caption: Optional[str]=None) -> str:
+        """Classic mode creates a prompt in a standard format first describing the image, 
+        then listing the artist, trending, movement, and flavor text modifiers."""
+        caption = caption or self.generate_caption(image)
         image_features = self.image_to_features(image)
 
         medium = self.mediums.rank(image_features, 1)[0]
@@ -158,73 +221,52 @@ class Interrogator():
 
         return _truncate_to_fit(prompt, self.tokenize)
 
-    def interrogate_fast(self, image: Image, max_flavors: int = 32) -> str:
-        caption = self.generate_caption(image)
+    def interrogate_fast(self, image: Image, max_flavors: int=32, caption: Optional[str]=None) -> str:
+        """Fast mode simply adds the top ranked terms after a caption. It generally results in 
+        better similarity between generated prompt and image than classic mode, but the prompts
+        are less readable."""
+        caption = caption or self.generate_caption(image)
         image_features = self.image_to_features(image)
-        merged = _merge_tables([self.artists, self.flavors, self.mediums, self.movements, self.trendings], self.config)
+        merged = _merge_tables([self.artists, self.flavors, self.mediums, self.movements, self.trendings], self)
         tops = merged.rank(image_features, max_flavors)
         return _truncate_to_fit(caption + ", " + ", ".join(tops), self.tokenize)
 
-    def interrogate(self, image: Image, max_flavors: int=32) -> str:
-        caption = self.generate_caption(image)
+    def interrogate_negative(self, image: Image, max_flavors: int = 32) -> str:
+        """Negative mode chains together the most dissimilar terms to the image. It can be used
+        to help build a negative prompt to pair with the regular positive prompt and often 
+        improve the results of generated images particularly with Stable Diffusion 2."""
         image_features = self.image_to_features(image)
+        flaves = self.flavors.rank(image_features, self.config.flavor_intermediate_count, reverse=True)
+        flaves = flaves + self.negative.labels
+        return self.chain(image_features, flaves, max_count=max_flavors, reverse=True, desc="Negative chain")
 
-        flaves = self.flavors.rank(image_features, self.config.flavor_intermediate_count)
-        best_medium = self.mediums.rank(image_features, 1)[0]
-        best_artist = self.artists.rank(image_features, 1)[0]
-        best_trending = self.trendings.rank(image_features, 1)[0]
-        best_movement = self.movements.rank(image_features, 1)[0]
-
-        best_prompt = caption
-        best_sim = self.similarity(image_features, best_prompt)
-
-        def check(addition: str) -> bool:
-            nonlocal best_prompt, best_sim
-            prompt = best_prompt + ", " + addition
-            sim = self.similarity(image_features, prompt)
-            if sim > best_sim:
-                best_sim = sim
-                best_prompt = prompt
-                return True
-            return False
-
-        def check_multi_batch(opts: List[str]):
-            nonlocal best_prompt, best_sim
-            prompts = []
-            for i in range(2**len(opts)):
-                prompt = best_prompt
-                for bit in range(len(opts)):
-                    if i & (1 << bit):
-                        prompt += ", " + opts[bit]
-                prompts.append(prompt)
-
-            t = LabelTable(prompts, None, self.clip_model, self.tokenize, self.config)
-            best_prompt = t.rank(image_features, 1)[0]
-            best_sim = self.similarity(image_features, best_prompt)
+    def interrogate(self, image: Image, min_flavors: int=8, max_flavors: int=32, caption: Optional[str]=None) -> str:
+        caption = caption or self.generate_caption(image)
+        image_features = self.image_to_features(image)
 
-        check_multi_batch([best_medium, best_artist, best_trending, best_movement])
+        merged = _merge_tables([self.artists, self.flavors, self.mediums, self.movements, self.trendings], self)
+        flaves = merged.rank(image_features, self.config.flavor_intermediate_count)
+        best_prompt, best_sim = caption, self.similarity(image_features, caption)
+        best_prompt = self.chain(image_features, flaves, best_prompt, best_sim, min_count=min_flavors, max_count=max_flavors, desc="Flavor chain")
 
-        extended_flavors = set(flaves)
-        for _ in tqdm(range(max_flavors), desc="Flavor chain", disable=self.config.quiet):
-            best = self.rank_top(image_features, [f"{best_prompt}, {f}" for f in extended_flavors])
-            flave = best[len(best_prompt)+2:]
-            if not check(flave):
-                break
-            if _prompt_at_max_len(best_prompt, self.tokenize):
-                break
-            extended_flavors.remove(flave)
+        fast_prompt = self.interrogate_fast(image, max_flavors, caption=caption)
+        classic_prompt = self.interrogate_classic(image, max_flavors, caption=caption)
+        candidates = [caption, classic_prompt, fast_prompt, best_prompt]
+        return candidates[np.argmax(self.similarities(image_features, candidates))]
 
-        return best_prompt
-
-    def rank_top(self, image_features: torch.Tensor, text_array: List[str]) -> str:
+    def rank_top(self, image_features: torch.Tensor, text_array: List[str], reverse: bool=False) -> str:
+        self._prepare_clip()
         text_tokens = self.tokenize([text for text in text_array]).to(self.device)
         with torch.no_grad(), torch.cuda.amp.autocast():
             text_features = self.clip_model.encode_text(text_tokens)
             text_features /= text_features.norm(dim=-1, keepdim=True)
             similarity = text_features @ image_features.T
+            if reverse:
+                similarity = -similarity
         return text_array[similarity.argmax().item()]
 
     def similarity(self, image_features: torch.Tensor, text: str) -> float:
+        self._prepare_clip()
         text_tokens = self.tokenize([text]).to(self.device)
         with torch.no_grad(), torch.cuda.amp.autocast():
             text_features = self.clip_model.encode_text(text_tokens)
@@ -232,29 +274,45 @@ class Interrogator():
             similarity = text_features @ image_features.T
         return similarity[0][0].item()
 
+    def similarities(self, image_features: torch.Tensor, text_array: List[str]) -> List[float]:
+        self._prepare_clip()
+        text_tokens = self.tokenize([text for text in text_array]).to(self.device)
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            text_features = self.clip_model.encode_text(text_tokens)
+            text_features /= text_features.norm(dim=-1, keepdim=True)
+            similarity = text_features @ image_features.T
+        return similarity.T[0].tolist()
+
+    def _prepare_caption(self):
+        if self.config.clip_offload and not self.clip_offloaded:
+            self.clip_model = self.clip_model.to('cpu')
+            self.clip_offloaded = True
+        if self.caption_offloaded:
+            self.caption_model = self.caption_model.to(self.device)
+            self.caption_offloaded = False
+
+    def _prepare_clip(self):
+        if self.config.caption_offload and not self.caption_offloaded:
+            self.caption_model = self.caption_model.to('cpu')
+            self.caption_offloaded = True
+        if self.clip_offloaded:
+            self.clip_model = self.clip_model.to(self.device)
+            self.clip_offloaded = False
+
 
 class LabelTable():
-    def __init__(self, labels:List[str], desc:str, clip_model, tokenize, config: Config):
+    def __init__(self, labels:List[str], desc:str, ci: Interrogator):
+        clip_model, config = ci.clip_model, ci.config
         self.chunk_size = config.chunk_size
         self.config = config
         self.device = config.device
         self.embeds = []
         self.labels = labels
-        self.tokenize = tokenize
+        self.tokenize = ci.tokenize
 
         hash = hashlib.sha256(",".join(labels).encode()).hexdigest()
-
-        cache_filepath = None
-        if config.cache_path is not None and desc is not None:
-            os.makedirs(config.cache_path, exist_ok=True)
-            sanitized_name = config.clip_model_name.replace('/', '_').replace('@', '_')
-            cache_filepath = os.path.join(config.cache_path, f"{sanitized_name}_{desc}.pkl")
-            if desc is not None and os.path.exists(cache_filepath):
-                with open(cache_filepath, 'rb') as f:
-                    data = pickle.load(f)
-                    if data.get('hash') == hash:
-                        self.labels = data['labels']
-                        self.embeds = data['embeds']
+        sanitized_name = self.config.clip_model_name.replace('/', '_').replace('@', '_')
+        self._load_cached(desc, hash, sanitized_name)
 
         if len(self.labels) != len(self.embeds):
             self.embeds = []
@@ -268,26 +326,63 @@ class LabelTable():
                 for i in range(text_features.shape[0]):
                     self.embeds.append(text_features[i])
 
-            if cache_filepath is not None:
-                with open(cache_filepath, 'wb') as f:
-                    pickle.dump({
-                        "labels": self.labels, 
-                        "embeds": self.embeds, 
-                        "hash": hash, 
-                        "model": config.clip_model_name
-                    }, f)
+            if desc and self.config.cache_path:
+                os.makedirs(self.config.cache_path, exist_ok=True)
+                cache_filepath = os.path.join(self.config.cache_path, f"{sanitized_name}_{desc}.safetensors")
+                tensors = {
+                    "embeds": np.stack(self.embeds),
+                    "hash": np.array([ord(c) for c in hash], dtype=np.int8)
+                }
+                save_file(tensors, cache_filepath)
+
+        if self.device == 'cpu' or self.device == torch.device('cpu'):
+            self.embeds = [e.astype(np.float32) for e in self.embeds]
+
+    def _load_cached(self, desc:str, hash:str, sanitized_name:str) -> bool:
+        if self.config.cache_path is None or desc is None:
+            return False
+
+        cached_safetensors = os.path.join(self.config.cache_path, f"{sanitized_name}_{desc}.safetensors")
+
+        if self.config.download_cache and not os.path.exists(cached_safetensors):
+            download_url = CACHE_URL_BASE + f"{sanitized_name}_{desc}.safetensors"
+            try:
+                os.makedirs(self.config.cache_path, exist_ok=True)
+                _download_file(download_url, cached_safetensors, quiet=self.config.quiet)
+            except Exception as e:
+                print(f"Failed to download {download_url}")
+                print(e)
+                return False                
+
+        if os.path.exists(cached_safetensors):
+            try:
+                tensors = load_file(cached_safetensors)
+            except Exception as e:
+                print(f"Failed to load {cached_safetensors}")
+                print(e)
+                return False
+            if 'hash' in tensors and 'embeds' in tensors:
+                if np.array_equal(tensors['hash'], np.array([ord(c) for c in hash], dtype=np.int8)):
+                    self.embeds = tensors['embeds']
+                    if len(self.embeds.shape) == 2:
+                        self.embeds = [self.embeds[i] for i in range(self.embeds.shape[0])]
+                    return True
+
+        return False
     
-    def _rank(self, image_features: torch.Tensor, text_embeds: torch.Tensor, top_count: int=1) -> str:
+    def _rank(self, image_features: torch.Tensor, text_embeds: torch.Tensor, top_count: int=1, reverse: bool=False) -> str:
         top_count = min(top_count, len(text_embeds))
         text_embeds = torch.stack([torch.from_numpy(t) for t in text_embeds]).to(self.device)
         with torch.cuda.amp.autocast():
             similarity = image_features @ text_embeds.T
+            if reverse:
+                similarity = -similarity
         _, top_labels = similarity.float().cpu().topk(top_count, dim=-1)
         return [top_labels[0][i].numpy() for i in range(top_count)]
 
-    def rank(self, image_features: torch.Tensor, top_count: int=1) -> List[str]:
+    def rank(self, image_features: torch.Tensor, top_count: int=1, reverse: bool=False) -> List[str]:
         if len(self.labels) <= self.chunk_size:
-            tops = self._rank(image_features, self.embeds, top_count=top_count)
+            tops = self._rank(image_features, self.embeds, top_count=top_count, reverse=reverse)
             return [self.labels[i] for i in tops]
 
         num_chunks = int(math.ceil(len(self.labels)/self.chunk_size))
@@ -297,7 +392,7 @@ class LabelTable():
         for chunk_idx in tqdm(range(num_chunks), disable=self.config.quiet):
             start = chunk_idx*self.chunk_size
             stop = min(start+self.chunk_size, len(self.embeds))
-            tops = self._rank(image_features, self.embeds[start:stop], top_count=keep_per_chunk)
+            tops = self._rank(image_features, self.embeds[start:stop], top_count=keep_per_chunk, reverse=reverse)
             top_labels.extend([self.labels[start+i] for i in tops])
             top_embeds.extend([self.embeds[start+i] for i in tops])
 
@@ -305,13 +400,23 @@ class LabelTable():
         return [top_labels[i] for i in tops]
 
 
-def _load_list(data_path: str, filename: str) -> List[str]:
-    with open(os.path.join(data_path, filename), 'r', encoding='utf-8', errors='replace') as f:
-        items = [line.strip() for line in f.readlines()]
-    return items
+def _download_file(url: str, filepath: str, chunk_size: int = 4*1024*1024, quiet: bool = False):
+    r = requests.get(url, stream=True)
+    if r.status_code != 200:
+        return
 
-def _merge_tables(tables: List[LabelTable], config: Config) -> LabelTable:
-    m = LabelTable([], None, None, None, config)
+    file_size = int(r.headers.get("Content-Length", 0))
+    filename = url.split("/")[-1]
+    progress = tqdm(total=file_size, unit="B", unit_scale=True, desc=filename, disable=quiet)
+    with open(filepath, "wb") as f:
+        for chunk in r.iter_content(chunk_size=chunk_size):
+            if chunk:
+                f.write(chunk)
+                progress.update(len(chunk))
+    progress.close()
+
+def _merge_tables(tables: List[LabelTable], ci: Interrogator) -> LabelTable:
+    m = LabelTable([], None, ci)
     for table in tables:
         m.labels.extend(table.labels)
         m.embeds.extend(table.embeds)
@@ -329,3 +434,17 @@ def _truncate_to_fit(text: str, tokenize) -> str:
             break
         new_text += ', ' + part
     return new_text
+
+def list_caption_models() -> List[str]:
+    return list(CAPTION_MODELS.keys())
+
+def list_clip_models() -> List[str]:
+    return ['/'.join(x) for x in open_clip.list_pretrained()]
+
+def load_list(data_path: str, filename: Optional[str] = None) -> List[str]:
+    """Load a list of strings from a file."""
+    if filename is not None:
+        data_path = os.path.join(data_path, filename)
+    with open(data_path, 'r', encoding='utf-8', errors='replace') as f:
+        items = [line.strip() for line in f.readlines()]
+    return items

clip_interrogator/data/negative.txt

@@ -0,0 +1,41 @@
+3d
+b&w
+bad anatomy
+bad art
+blur
+blurry
+cartoon
+childish
+close up
+deformed
+disconnected limbs
+disfigured
+disgusting
+extra limb
+extra limbs
+floating limbs
+grain
+illustration
+kitsch
+long body
+long neck
+low quality
+low-res
+malformed hands
+mangled
+missing limb
+mutated
+mutation
+mutilated
+noisy
+old
+out of focus
+over saturation
+oversaturated
+poorly drawn
+poorly drawn face
+poorly drawn hands
+render
+surreal
+ugly
+weird colors

cog.yaml

@@ -1,21 +1,16 @@
 build:
   gpu: true
-  cuda: "11.6"
-  python_version: "3.8"
+  cuda: "11.8"
+  python_version: "3.10"
   system_packages:
     - "libgl1-mesa-glx"
     - "libglib2.0-0"
   python_packages:
-    - "ipython==8.4.0"
-    - "fairscale==0.4.12"
-    - "transformers==4.21.2"
-    - "ftfy==6.1.1"
-    - "torch==1.13.0 --extra-index-url=https://download.pytorch.org/whl/cu116"
-    - "torchvision==0.14.0 --extra-index-url=https://download.pytorch.org/whl/cu116"
-    - "open_clip_torch==2.7.0"    
-    - "timm==0.4.12"
-    - "pycocoevalcap==1.2"
-  run:
-  - git clone https://github.com/salesforce/BLIP /root/blip
+    - "Pillow==10.0.0"
+    - "safetensors==0.3.3"
+    - "tqdm==4.66.1"
+    - "open_clip_torch==2.20.0"
+    - "accelerate==0.22.0"
+    - "transformers==4.33.1"
 
 predict: "predict.py:Predictor"

predict.py

@@ -2,15 +2,12 @@ import sys
 from PIL import Image
 from cog import BasePredictor, Input, Path
 
-sys.path.append('/root/blip')
-
-from clip_interrogator import Interrogator, Config
+from clip_interrogator import Config, Interrogator
 
 
 class Predictor(BasePredictor):
     def setup(self):
         self.ci = Interrogator(Config(
-            blip_model_url='cache/model_large_caption.pth',
             clip_model_name="ViT-L-14/openai",
             clip_model_path='cache',
             device='cuda:0', 
@@ -21,18 +18,27 @@ class Predictor(BasePredictor):
         image: Path = Input(description="Input image"),
         clip_model_name: str = Input(
             default="ViT-L-14/openai",
-            choices=[
-                "ViT-L-14/openai",
-                "ViT-H-14/laion2b_s32b_b79k",
-            ],
-            description="Choose a clip model.",
+            choices=["ViT-L-14/openai", "ViT-H-14/laion2b_s32b_b79k", "ViT-bigG-14/laion2b_s39b_b160k"],
+            description="Choose ViT-L for Stable Diffusion 1, ViT-H for Stable Diffusion 2, or ViT-bigG for Stable Diffusion XL.",
+        ),        
+        mode: str = Input(
+            default="best",
+            choices=["best", "classic", "fast", "negative"],
+            description="Prompt mode (best takes 10-20 seconds, fast takes 1-2 seconds).",
         ),        
     ) -> str:
         """Run a single prediction on the model"""
         image = Image.open(str(image)).convert("RGB")
         self.switch_model(clip_model_name)
-        return self.ci.interrogate(image)
-    
+        if mode == 'best':
+            return self.ci.interrogate(image)
+        elif mode == 'classic':
+            return self.ci.interrogate_classic(image)
+        elif mode == 'fast':
+            return self.ci.interrogate_fast(image)
+        elif mode == 'negative':
+            return self.ci.interrogate_negative(image)
+
     def switch_model(self, clip_model_name: str):
         if clip_model_name != self.ci.config.clip_model_name:
             self.ci.config.clip_model_name = clip_model_name

requirements.txt

@@ -1,6 +1,9 @@
-torch
+torch>=1.13.0
 torchvision
 Pillow
 requests
+safetensors
 tqdm
-open_clip_torch
+open_clip_torch
+accelerate
+transformers>=4.27.1

run_cli.py

@@ -1,12 +1,11 @@
 #!/usr/bin/env python3
 import argparse
 import csv
-import open_clip
 import os
 import requests
 import torch
 from PIL import Image
-from clip_interrogator import Interrogator, Config
+from clip_interrogator import Interrogator, Config, list_clip_models
 
 def inference(ci, image, mode):
     image = image.convert('RGB')
@@ -20,9 +19,11 @@ def inference(ci, image, mode):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('-c', '--clip', default='ViT-L-14/openai', help='name of CLIP model to use')
+    parser.add_argument('-d', '--device', default='auto', help='device to use (auto, cuda or cpu)')
     parser.add_argument('-f', '--folder', help='path to folder of images')
     parser.add_argument('-i', '--image', help='image file or url')
     parser.add_argument('-m', '--mode', default='best', help='best, classic, or fast')
+    parser.add_argument("--lowvram", action='store_true', help="Optimize settings for low VRAM")
 
     args = parser.parse_args()
     if not args.folder and not args.image:
@@ -34,15 +35,24 @@ def main():
         exit(1)
 
     # validate clip model name
-    models = ['/'.join(x) for x in open_clip.list_pretrained()]
+    models = list_clip_models()
     if args.clip not in models:
         print(f"Could not find CLIP model {args.clip}!")
         print(f"    available models: {models}")
         exit(1)
 
+    # select device
+    if args.device == 'auto':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        if not torch.cuda.is_available():
+            print("CUDA is not available, using CPU. Warning: this will be very slow!")
+    else:
+        device = torch.device(args.device)
+
     # generate a nice prompt
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     config = Config(device=device, clip_model_name=args.clip)
+    if args.lowvram:
+        config.apply_low_vram_defaults()
     ci = Interrogator(config)
 
     # process single image

run_gradio.py

@@ -1,49 +1,99 @@
 #!/usr/bin/env python3
 import argparse
-import gradio as gr
-import open_clip
-from clip_interrogator import Interrogator, Config
+import torch
+from clip_interrogator import Config, Interrogator, list_caption_models, list_clip_models
+
+try:
+    import gradio as gr
+except ImportError:
+    print("Gradio is not installed, please install it with 'pip install gradio'")
+    exit(1)
 
 parser = argparse.ArgumentParser()
+parser.add_argument("--lowvram", action='store_true', help="Optimize settings for low VRAM")
 parser.add_argument('-s', '--share', action='store_true', help='Create a public link')
 args = parser.parse_args()
 
-ci = Interrogator(Config(cache_path="cache", clip_model_path="cache"))
+if not torch.cuda.is_available():
+    print("CUDA is not available, using CPU. Warning: this will be very slow!")
+
+config = Config(cache_path="cache")
+if args.lowvram:
+    config.apply_low_vram_defaults()
+ci = Interrogator(config)
+
+def image_analysis(image, clip_model_name):
+    if clip_model_name != ci.config.clip_model_name:
+        ci.config.clip_model_name = clip_model_name
+        ci.load_clip_model()
+
+    image = image.convert('RGB')
+    image_features = ci.image_to_features(image)
+
+    top_mediums = ci.mediums.rank(image_features, 5)
+    top_artists = ci.artists.rank(image_features, 5)
+    top_movements = ci.movements.rank(image_features, 5)
+    top_trendings = ci.trendings.rank(image_features, 5)
+    top_flavors = ci.flavors.rank(image_features, 5)
+
+    medium_ranks = {medium: sim for medium, sim in zip(top_mediums, ci.similarities(image_features, top_mediums))}
+    artist_ranks = {artist: sim for artist, sim in zip(top_artists, ci.similarities(image_features, top_artists))}
+    movement_ranks = {movement: sim for movement, sim in zip(top_movements, ci.similarities(image_features, top_movements))}
+    trending_ranks = {trending: sim for trending, sim in zip(top_trendings, ci.similarities(image_features, top_trendings))}
+    flavor_ranks = {flavor: sim for flavor, sim in zip(top_flavors, ci.similarities(image_features, top_flavors))}
+    
+    return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
+
+def image_to_prompt(image, mode, clip_model_name, blip_model_name):
+    if blip_model_name != ci.config.caption_model_name:
+        ci.config.caption_model_name = blip_model_name
+        ci.load_caption_model()
 
-def inference(image, mode, clip_model_name, blip_max_length, blip_num_beams):
     if clip_model_name != ci.config.clip_model_name:
         ci.config.clip_model_name = clip_model_name
         ci.load_clip_model()
-    ci.config.blip_max_length = int(blip_max_length)
-    ci.config.blip_num_beams = int(blip_num_beams)
 
     image = image.convert('RGB')
     if mode == 'best':
         return ci.interrogate(image)
     elif mode == 'classic':
         return ci.interrogate_classic(image)
-    else:
+    elif mode == 'fast':
         return ci.interrogate_fast(image)
+    elif mode == 'negative':
+        return ci.interrogate_negative(image)
+
+def prompt_tab():
+    with gr.Column():
+        with gr.Row():
+            image = gr.Image(type='pil', label="Image")
+            with gr.Column():
+                mode = gr.Radio(['best', 'fast', 'classic', 'negative'], label='Mode', value='best')
+                clip_model = gr.Dropdown(list_clip_models(), value=ci.config.clip_model_name, label='CLIP Model')
+                blip_model = gr.Dropdown(list_caption_models(), value=ci.config.caption_model_name, label='Caption Model')
+        prompt = gr.Textbox(label="Prompt")
+    button = gr.Button("Generate prompt")
+    button.click(image_to_prompt, inputs=[image, mode, clip_model, blip_model], outputs=prompt)
+
+def analyze_tab():
+    with gr.Column():
+        with gr.Row():
+            image = gr.Image(type='pil', label="Image")
+            model = gr.Dropdown(list_clip_models(), value='ViT-L-14/openai', label='CLIP Model')
+        with gr.Row():
+            medium = gr.Label(label="Medium", num_top_classes=5)
+            artist = gr.Label(label="Artist", num_top_classes=5)        
+            movement = gr.Label(label="Movement", num_top_classes=5)
+            trending = gr.Label(label="Trending", num_top_classes=5)
+            flavor = gr.Label(label="Flavor", num_top_classes=5)
+    button = gr.Button("Analyze")
+    button.click(image_analysis, inputs=[image, model], outputs=[medium, artist, movement, trending, flavor])
 
-models = ['/'.join(x) for x in open_clip.list_pretrained()]
-
-inputs = [
-    gr.inputs.Image(type='pil'),
-    gr.Radio(['best', 'classic', 'fast'], label='Mode', value='best'),
-    gr.Dropdown(models, value='ViT-L-14/openai', label='CLIP Model'),
-    gr.Number(value=32, label='Caption Max Length'),
-    gr.Number(value=64, label='Caption Num Beams'),
-]
-outputs = [
-    gr.outputs.Textbox(label="Output"),
-]
-
-io = gr.Interface(
-    inference, 
-    inputs, 
-    outputs, 
-    title="🕵️‍♂️ CLIP Interrogator 🕵️‍♂️",
-    allow_flagging=False,
-)
-io.launch(share=args.share)
+with gr.Blocks() as ui:
+    gr.Markdown("# <center>🕵️‍♂️ CLIP Interrogator 🕵️‍♂️</center>")
+    with gr.Tab("Prompt"):
+        prompt_tab()
+    with gr.Tab("Analyze"):
+        analyze_tab()
 
+ui.launch(show_api=False, debug=True, share=args.share)

setup.py

@@ -5,13 +5,13 @@ from setuptools import setup, find_packages
 
 setup(
     name="clip-interrogator",
-    version="0.2.0",
+    version="0.6.0",
     license='MIT',
     author='pharmapsychotic',
     author_email='me@pharmapsychotic.com',
     url='https://github.com/pharmapsychotic/clip-interrogator',
     description="Generate a prompt from an image",
-    long_description=open('README.md').read(),
+    long_description=open('README.md', encoding='utf-8').read(),
     long_description_content_type="text/markdown",
     packages=find_packages(),
     install_requires=[