copy class dependencies to notebook

Can't figure out how to run it in the cloud with local deps yet

Cliptag.ipynb

@@ -2,7 +2,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -10,10 +10,156 @@
         "id": "WDAOVA8dtmR-",
         "outputId": "ca96e5b3-5566-4264-d4e4-71b1e7ef8c49"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Requirement already satisfied: torch in ./.venv/lib/python3.11/site-packages (2.0.0)\n",
+            "Requirement already satisfied: filelock in ./.venv/lib/python3.11/site-packages (from torch) (3.12.0)\n",
+            "Requirement already satisfied: typing-extensions in ./.venv/lib/python3.11/site-packages (from torch) (4.5.0)\n",
+            "Requirement already satisfied: sympy in ./.venv/lib/python3.11/site-packages (from torch) (1.11.1)\n",
+            "Requirement already satisfied: networkx in ./.venv/lib/python3.11/site-packages (from torch) (3.1)\n",
+            "Requirement already satisfied: jinja2 in ./.venv/lib/python3.11/site-packages (from torch) (3.1.2)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in ./.venv/lib/python3.11/site-packages (from jinja2->torch) (2.1.2)\n",
+            "Requirement already satisfied: mpmath>=0.19 in ./.venv/lib/python3.11/site-packages (from sympy->torch) (1.3.0)\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "Note: you may need to restart the kernel to use updated packages.\n",
+            "Requirement already satisfied: clip_interrogator in ./.venv/lib/python3.11/site-packages (0.6.0)\n",
+            "Requirement already satisfied: torch in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (2.0.0)\n",
+            "Requirement already satisfied: torchvision in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (0.15.1)\n",
+            "Requirement already satisfied: Pillow in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (9.5.0)\n",
+            "Requirement already satisfied: requests in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (2.28.2)\n",
+            "Requirement already satisfied: safetensors in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (0.3.1)\n",
+            "Requirement already satisfied: tqdm in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (4.65.0)\n",
+            "Requirement already satisfied: open-clip-torch in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (2.17.1)\n",
+            "Requirement already satisfied: accelerate in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (0.18.0)\n",
+            "Requirement already satisfied: transformers>=4.27.1 in ./.venv/lib/python3.11/site-packages (from clip_interrogator) (4.28.1)\n",
+            "Requirement already satisfied: filelock in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (3.12.0)\n",
+            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (0.13.4)\n",
+            "Requirement already satisfied: numpy>=1.17 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (1.24.2)\n",
+            "Requirement already satisfied: packaging>=20.0 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (23.1)\n",
+            "Requirement already satisfied: pyyaml>=5.1 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (6.0)\n",
+            "Requirement already satisfied: regex!=2019.12.17 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (2023.3.23)\n",
+            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in ./.venv/lib/python3.11/site-packages (from transformers>=4.27.1->clip_interrogator) (0.13.3)\n",
+            "Requirement already satisfied: psutil in ./.venv/lib/python3.11/site-packages (from accelerate->clip_interrogator) (5.9.5)\n",
+            "Requirement already satisfied: typing-extensions in ./.venv/lib/python3.11/site-packages (from torch->clip_interrogator) (4.5.0)\n",
+            "Requirement already satisfied: sympy in ./.venv/lib/python3.11/site-packages (from torch->clip_interrogator) (1.11.1)\n",
+            "Requirement already satisfied: networkx in ./.venv/lib/python3.11/site-packages (from torch->clip_interrogator) (3.1)\n",
+            "Requirement already satisfied: jinja2 in ./.venv/lib/python3.11/site-packages (from torch->clip_interrogator) (3.1.2)\n",
+            "Requirement already satisfied: ftfy in ./.venv/lib/python3.11/site-packages (from open-clip-torch->clip_interrogator) (6.1.1)\n",
+            "Requirement already satisfied: sentencepiece in ./.venv/lib/python3.11/site-packages (from open-clip-torch->clip_interrogator) (0.1.98)\n",
+            "Requirement already satisfied: protobuf<4 in ./.venv/lib/python3.11/site-packages (from open-clip-torch->clip_interrogator) (3.20.3)\n",
+            "Requirement already satisfied: timm in ./.venv/lib/python3.11/site-packages (from open-clip-torch->clip_interrogator) (0.6.13)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in ./.venv/lib/python3.11/site-packages (from requests->clip_interrogator) (3.1.0)\n",
+            "Requirement already satisfied: idna<4,>=2.5 in ./.venv/lib/python3.11/site-packages (from requests->clip_interrogator) (3.4)\n",
+            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in ./.venv/lib/python3.11/site-packages (from requests->clip_interrogator) (1.26.15)\n",
+            "Requirement already satisfied: certifi>=2017.4.17 in ./.venv/lib/python3.11/site-packages (from requests->clip_interrogator) (2022.12.7)\n",
+            "Requirement already satisfied: wcwidth>=0.2.5 in ./.venv/lib/python3.11/site-packages (from ftfy->open-clip-torch->clip_interrogator) (0.2.6)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in ./.venv/lib/python3.11/site-packages (from jinja2->torch->clip_interrogator) (2.1.2)\n",
+            "Requirement already satisfied: mpmath>=0.19 in ./.venv/lib/python3.11/site-packages (from sympy->torch->clip_interrogator) (1.3.0)\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "Note: you may need to restart the kernel to use updated packages.\n",
+            "Requirement already satisfied: open_clip_torch in ./.venv/lib/python3.11/site-packages (2.17.1)\n",
+            "Requirement already satisfied: torch>=1.9.0 in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (2.0.0)\n",
+            "Requirement already satisfied: torchvision in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (0.15.1)\n",
+            "Requirement already satisfied: regex in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (2023.3.23)\n",
+            "Requirement already satisfied: ftfy in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (6.1.1)\n",
+            "Requirement already satisfied: tqdm in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (4.65.0)\n",
+            "Requirement already satisfied: huggingface-hub in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (0.13.4)\n",
+            "Requirement already satisfied: sentencepiece in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (0.1.98)\n",
+            "Requirement already satisfied: protobuf<4 in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (3.20.3)\n",
+            "Requirement already satisfied: timm in ./.venv/lib/python3.11/site-packages (from open_clip_torch) (0.6.13)\n",
+            "Requirement already satisfied: filelock in ./.venv/lib/python3.11/site-packages (from torch>=1.9.0->open_clip_torch) (3.12.0)\n",
+            "Requirement already satisfied: typing-extensions in ./.venv/lib/python3.11/site-packages (from torch>=1.9.0->open_clip_torch) (4.5.0)\n",
+            "Requirement already satisfied: sympy in ./.venv/lib/python3.11/site-packages (from torch>=1.9.0->open_clip_torch) (1.11.1)\n",
+            "Requirement already satisfied: networkx in ./.venv/lib/python3.11/site-packages (from torch>=1.9.0->open_clip_torch) (3.1)\n",
+            "Requirement already satisfied: jinja2 in ./.venv/lib/python3.11/site-packages (from torch>=1.9.0->open_clip_torch) (3.1.2)\n",
+            "Requirement already satisfied: wcwidth>=0.2.5 in ./.venv/lib/python3.11/site-packages (from ftfy->open_clip_torch) (0.2.6)\n",
+            "Requirement already satisfied: requests in ./.venv/lib/python3.11/site-packages (from huggingface-hub->open_clip_torch) (2.28.2)\n",
+            "Requirement already satisfied: pyyaml>=5.1 in ./.venv/lib/python3.11/site-packages (from huggingface-hub->open_clip_torch) (6.0)\n",
+            "Requirement already satisfied: packaging>=20.9 in ./.venv/lib/python3.11/site-packages (from huggingface-hub->open_clip_torch) (23.1)\n",
+            "Requirement already satisfied: numpy in ./.venv/lib/python3.11/site-packages (from torchvision->open_clip_torch) (1.24.2)\n",
+            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in ./.venv/lib/python3.11/site-packages (from torchvision->open_clip_torch) (9.5.0)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in ./.venv/lib/python3.11/site-packages (from jinja2->torch>=1.9.0->open_clip_torch) (2.1.2)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in ./.venv/lib/python3.11/site-packages (from requests->huggingface-hub->open_clip_torch) (3.1.0)\n",
+            "Requirement already satisfied: idna<4,>=2.5 in ./.venv/lib/python3.11/site-packages (from requests->huggingface-hub->open_clip_torch) (3.4)\n",
+            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in ./.venv/lib/python3.11/site-packages (from requests->huggingface-hub->open_clip_torch) (1.26.15)\n",
+            "Requirement already satisfied: certifi>=2017.4.17 in ./.venv/lib/python3.11/site-packages (from requests->huggingface-hub->open_clip_torch) (2022.12.7)\n",
+            "Requirement already satisfied: mpmath>=0.19 in ./.venv/lib/python3.11/site-packages (from sympy->torch>=1.9.0->open_clip_torch) (1.3.0)\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "Note: you may need to restart the kernel to use updated packages.\n"
+          ]
+        }
+      ],
       "source": [
         "%pip install torch\n",
-        "%pip install clip_interrogator"
+        "%pip install clip_interrogator\n",
+        "%pip install open_clip_torch"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "from clip_interrogator import Interrogator, LabelTable, Config\n",
+        "import open_clip\n",
+        "import os\n",
+        "\n",
+        "def load_keywords_from_files(directory):\n",
+        "    keyword_list = []\n",
+        "    for filename in os.listdir(directory):\n",
+        "        # Check if the file has a .txt extension\n",
+        "        if filename.endswith('.txt'):\n",
+        "            file_path = os.path.join(directory, filename)\n",
+        "            with open(file_path, 'r') as file:\n",
+        "                # Read lines and strip leading/trailing whitespace\n",
+        "                keywords = [line.strip() for line in file]\n",
+        "                # Add keywords to the master keyword_list\n",
+        "                keyword_list.extend(keywords)\n",
+        "    print(f\"Loaded {len(keyword_list)} keywords...\")\n",
+        "    return keyword_list\n",
+        "\n",
+        "class CIInterrogator(Interrogator):\n",
+        "    def __init__(self, config: Config):\n",
+        "        self.config = config\n",
+        "        self.device = config.device\n",
+        "        self.dtype = torch.float16 if self.device == 'cuda' else torch.float32\n",
+        "        self.caption_offloaded = True\n",
+        "        self.clip_offloaded = True\n",
+        "        self.load_clip_model()\n",
+        "\n",
+        "    def load_clip_model(self):\n",
+        "        config = self.config\n",
+        "\n",
+        "        clip_model_name, clip_model_pretrained_name = config.clip_model_name.split('/', 2)\n",
+        "\n",
+        "        if config.clip_model is None:\n",
+        "            if not config.quiet:\n",
+        "                print(f\"Loading CLIP model {config.clip_model_name}...\")\n",
+        "\n",
+        "            self.clip_model, _, self.clip_preprocess = open_clip.create_model_and_transforms(\n",
+        "                clip_model_name, \n",
+        "                pretrained=clip_model_pretrained_name, \n",
+        "                precision='fp16' if config.device == 'cuda' else 'fp32',\n",
+        "                device=config.device,\n",
+        "                jit=False,\n",
+        "                cache_dir=config.clip_model_path\n",
+        "            )\n",
+        "            self.clip_model.eval()\n",
+        "        else:\n",
+        "            self.clip_model = config.clip_model\n",
+        "            self.clip_preprocess = config.clip_preprocess\n",
+        "        self.tokenize = open_clip.get_tokenizer(clip_model_name)\n"
       ]
     },
     {
@@ -33,7 +179,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": null,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -41,20 +187,9 @@
         "id": "bJ1tJyL0-8qQ",
         "outputId": "1313dc44-baba-4d49-c1ee-ebf6b0c351f6"
       },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Using CPU\n",
-            "Loading CLIP model ViT-L-14/openai...\n",
-            "Loading CLIP model ViT-B-32/laion2b_s34b_b79k...\n"
-          ]
-        }
-      ],
+      "outputs": [],
       "source": [
         "from clip_interrogator import Config\n",
-        "from cliptag.cliptag import load_keywords_from_files, CIInterrogator\n",
         "import torch \n",
         "\n",
         "if torch.cuda.is_available():\n",