From bc07ce62c179d3aab3053a623d96a071101d11cb Mon Sep 17 00:00:00 2001
From: pharmapsychotic <me@pharmapsychotic.com>
Date: Sat, 9 Sep 2023 19:11:06 -0500
Subject: [PATCH] Update cog

---
 cog.yaml   | 20 ++++++++------------
 predict.py | 19 +++++++++++--------
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/cog.yaml b/cog.yaml
index 638f2c6..da14215 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -1,20 +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"
-    - "git+https://github.com/pharmapsychotic/BLIP.git"
+    - "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"
diff --git a/predict.py b/predict.py
index abceb28..19625f1 100644
--- a/predict.py
+++ b/predict.py
@@ -2,13 +2,12 @@ import sys
 from PIL import Image
 from cog import BasePredictor, Input, Path
 
-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', 
@@ -19,23 +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 ViT-L for Stable Diffusion 1, and ViT-H for Stable Diffusion 2",
+            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", "fast"],
+            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)
-        if mode == "best":
+        if mode == 'best':
             return self.ci.interrogate(image)
-        else:
+        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