lazily steal the other classifiers from the autowebui extension

__init__.py

@@ -13,20 +13,22 @@ import folder_paths
 import model_management
 import nodes
 
-AESTHETIC_MODELS = {"aesthetic": 2, "waifu": 5}
+import sys
+sys.path.append(str(pathlib.Path(__file__).parent))
+
+CAFE_MODELS = {"cafe_aesthetic": 2, "cafe_waifu": 5}
 BLOCK_ORDER = [12, 11, 13, 10, 14, 9, 15, 8, 16, 7, 17, 6, 18,
                5, 19, 4, 20, 3, 21, 2, 22, 1, 23, 0, 24]
 
 
-def run_classifier(tensors, classifier):
-    image = Image.fromarray(np.clip(255. * tensors.cpu().numpy().squeeze(),
-                                    0, 255).astype(np.uint8))
+def run_cafe_classifier(image, classifier):
+    import transformers
     pipe = transformers.pipeline(
         "image-classification",
-        model=f"cafeai/cafe_{classifier}")
-    result = pipe(image, top_k=AESTHETIC_MODELS[classifier])
+        model=f"cafeai/{classifier}")
+    result = pipe(image, top_k=CAFE_MODELS[classifier])
     for data in result:
-        if data['label'] == classifier:
+        if data['label'] == classifier.split("_")[1]:
             return data['score']
 
 
@@ -52,7 +54,7 @@ class AutoMBW:
                 }),
                 "search_depth": ("INT", {"default": 4, "min": 2}),
                 "sample_count": ("INT", {"default": 1, "min": 1}),
-                "classifier": (["aesthetic", "waifu"],),
+                "classifier": (["aesthetic", "laion", "cafe_aesthetic", "cafe_waifu"],),
             }}
 
     RETURN_TYPES = ()
@@ -77,19 +79,28 @@ class AutoMBW:
             sd1[key].copy_(self.blocks_backup[key])
 
     def rate_model(self):
-        score = 0
+        rating = 0
         for i in range(self.sample_count):
             latent = nodes.common_ksampler(
                 self.model1, i, 20, 7.0, "ddim", "normal", self.prompt,
                 self.negative, {"samples": torch.zeros([1, 4, 64, 64])},
                 denoise=1.0)
-            image = self.vae.decode(latent[0]["samples"])
+            decoded = self.vae.decode(latent[0]["samples"])
+            image = Image.fromarray(
+                np.clip(255. * decoded.cpu().numpy().squeeze(),
+                        0, 255).astype(np.uint8))
             with warnings.catch_warnings():
                 # several possible transformers nags
                 warnings.filterwarnings('ignore')
-                score += run_classifier(
-                    image, self.classifier) / self.sample_count
-        return score
+                if self.classifier.startswith("cafe_"):
+                    rating += run_cafe_classifier(image, self.classifier)
+                elif self.classifier == "laion":
+                    import laion.score_laion_sac_logos_ava_v2
+                    rating += laion.score_laion_sac_logos_ava_v2.score(image)
+                elif self.classifier == "aesthetic":
+                    import aesthetic.score_aes_B32_v0
+                    rating += aesthetic.score_aes_B32_v0.score(image)
+        return rating
 
     def search(self, block, current, start, depth, maximum):
         if depth > self.search_depth or current > 1 or current < 0:

aesthetic/aesthetic.py

@@ -0,0 +1,30 @@
+import torch
+import numpy as np
+
+use_cuda = torch.cuda.is_available()
+
+def image_embeddings_direct(image, model, processor):
+    inputs = processor(images=image, return_tensors='pt')['pixel_values']
+    if use_cuda:
+        inputs = inputs.to('cuda')
+    result = model.get_image_features(pixel_values=inputs).cpu().detach().numpy()
+    return (result / np.linalg.norm(result)).squeeze(axis=0)
+
+# binary classifier that consumes CLIP embeddings
+class Classifier(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
+        self.fc2 = torch.nn.Linear(hidden_size, hidden_size//2)
+        self.fc3 = torch.nn.Linear(hidden_size//2, output_size)
+        self.relu = torch.nn.ReLU()
+        self.sigmoid = torch.nn.Sigmoid()
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.relu(x)
+        x = self.fc2(x)
+        x = self.relu(x)
+        x = self.fc3(x)
+        x = self.sigmoid(x)
+        return x

aesthetic/score_aes_B32_v0.py

@@ -0,0 +1,19 @@
+import os
+import torch
+import safetensors
+from transformers import CLIPModel, CLIPProcessor
+from aesthetic.aesthetic import image_embeddings_direct, Classifier
+
+dirname = os.path.dirname(__file__)
+aesthetic_path = os.path.join(dirname, "aes-B32-v0.safetensors")
+clip_name = 'openai/clip-vit-base-patch32'
+clipprocessor = CLIPProcessor.from_pretrained(clip_name)
+clipmodel = CLIPModel.from_pretrained(clip_name).to('cuda').eval()
+aes_model = Classifier(512, 256, 1).to('cuda')
+aes_model.load_state_dict(safetensors.torch.load_file(aesthetic_path))
+
+def score(image):
+    image_embeds = image_embeddings_direct(image, clipmodel, clipprocessor)
+    prediction = aes_model(torch.from_numpy(image_embeds).float().to('cuda'))
+    return prediction.item()
+

laion/laion.py

@@ -0,0 +1,47 @@
+import torch
+import numpy as np
+import clip
+
+use_cuda = torch.cuda.is_available()
+
+def image_embeddings_direct(image, model, processor):
+    inputs = processor(images=image, return_tensors='pt')['pixel_values']
+    if use_cuda:
+        inputs = inputs.to('cuda')
+    result = model.get_image_features(pixel_values=inputs).cpu().detach().numpy()
+    return (result / np.linalg.norm(result)).squeeze(axis=0)
+
+def normalized(a, axis=-1, order=2):
+    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
+    l2[l2 == 0] = 1
+    return a / np.expand_dims(l2, axis)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model, preprocess = clip.load("ViT-L/14", device=device)
+
+def image_embeddings_direct_laion(pil_image):
+    image = preprocess(pil_image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        image_features = model.encode_image(image)
+    im_emb_arr = normalized(image_features.cpu().detach().numpy())
+    return im_emb_arr
+
+class MLP(torch.nn.Module):
+    def __init__(self, input_size, xcol='emb', ycol='avg_rating'):
+        super().__init__()
+        self.input_size = input_size
+        self.xcol = xcol
+        self.ycol = ycol
+        self.layers = torch.nn.Sequential(
+            torch.nn.Linear(self.input_size, 1024),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(1024, 128),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(128, 64),
+            torch.nn.Dropout(0.1),
+            torch.nn.Linear(64, 16),
+            torch.nn.Linear(16, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)

laion/score_laion_sac_logos_ava_v2.py

@@ -0,0 +1,15 @@
+import os
+import torch
+import safetensors
+from laion.laion import image_embeddings_direct_laion, MLP
+
+dirname = os.path.dirname(__file__)
+aesthetic_path = os.path.join(dirname, "laion-sac-logos-ava-v2.safetensors")
+aes_model = MLP(768).to('cuda').eval()
+aes_model.load_state_dict(safetensors.torch.load_file(aesthetic_path))
+
+def score(image):
+    image_embeds = image_embeddings_direct_laion(image)
+    prediction = aes_model(torch.from_numpy(image_embeds).float().to('cuda'))
+    return prediction.item()
+