Upload 8 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/1.png filter=lfs diff=lfs merge=lfs -text
+examples/2.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,20 @@
+import gradio as gr
+from modules.predict import WaifuScorer
+
+
+def ui():
+    scorer = WaifuScorer()
+
+    def predict(img):
+        return scorer(img)
+
+    interface = gr.Interface(title='Waifu Scorer', description='A model that scores an anime illustration (0 ~ 10).',
+                             fn=predict, inputs=gr.Image(sources='upload', type='pil', height=512), outputs=gr.Number(precision=3),
+                             allow_flagging='never', examples=[['./examples/1.png'], ['./examples/2.png']]
+                             )
+
+    return interface
+
+
+if __name__ == '__main__':
+    ui().launch()

modules/__init__.py

@@ -0,0 +1 @@
+from .predict import WaifuScorer

modules/mlp.py

@@ -0,0 +1,34 @@
+import torch.nn as nn
+import pytorch_lightning as pl
+
+
+class MLP(pl.LightningModule):
+    def __init__(self, input_size, xcol='emb', ycol='avg_rating', batch_norm=True):
+        super().__init__()
+        self.input_size = input_size
+        self.xcol = xcol
+        self.ycol = ycol
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 2048),
+            nn.ReLU(),
+            nn.BatchNorm1d(2048) if batch_norm else nn.Identity(),
+            nn.Dropout(0.3),
+            nn.Linear(2048, 512),
+            nn.ReLU(),
+            nn.BatchNorm1d(512) if batch_norm else nn.Identity(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.BatchNorm1d(256) if batch_norm else nn.Identity(),
+            nn.Dropout(0.2),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.BatchNorm1d(128) if batch_norm else nn.Identity(),
+            nn.Dropout(0.1),
+            nn.Linear(128, 32),
+            nn.ReLU(),
+            nn.Linear(32, 1)
+        )
+
+    def forward(self, x):
+        return self.layers(x)

modules/predict.py

@@ -0,0 +1,84 @@
+import torch
+import clip
+import os
+import time
+from PIL import Image
+from typing import List
+from .mlp import MLP
+from .utils import download_from_url
+
+MLP_MODEL_URL = "https://huggingface.co/Eugeoter/waifu-scorer/waifu-scorer-v1-large.pth"
+
+
+class WaifuScorer:
+    def __init__(self, model_path: str = None, device: str = 'cuda', verbose=False):
+        self.verbose = verbose
+
+        if self.verbose:
+            tic = time.time()
+            print(f"loading pretrained model from `{model_path}`")
+
+        if model_path is None or not os.path.isfile(model_path):
+            model_path = download_from_url(MLP_MODEL_URL)
+
+        if device == 'cuda' and not torch.cuda.is_available():
+            device = 'cpu'
+            print("CUDA is not available, using CPU instead")
+
+        self.mlp = load_model(model_path, input_size=768, device=device)
+        self.model2, self.preprocess = load_clip_models("ViT-L/14", device=device)
+        self.device = self.mlp.device
+        self.dtype = self.mlp.dtype
+
+        self.mlp.eval()
+
+        if self.verbose:
+            toc = time.time()
+            print(f"model loaded: time_cost={toc-tic:.2f} | device={self.device} | dtype={self.dtype}")
+
+    @torch.no_grad()
+    def __call__(self, images: List[Image.Image]) -> List[float]:
+        if isinstance(images, Image.Image):
+            images = [images]
+        n = len(images)
+        if n == 1:
+            images = images*2  # batch norm
+        images = encode_images(images, self.model2, self.preprocess, device=self.device).to(device=self.device, dtype=self.dtype)
+        predictions = self.mlp(images)
+        scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
+        if n == 1:
+            scores = scores[0]
+        return scores
+
+
+def load_clip_models(name: str = "ViT-L/14", device='cuda'):
+    model2, preprocess = clip.load(name, device=device)  # RN50x64
+    return model2, preprocess
+
+
+def load_model(model_path: str = None, input_size=768, device: str = 'cuda', dtype=None):
+    model = MLP(input_size=input_size)
+    if model_path:
+        s = torch.load(model_path, map_location=device)
+        model.load_state_dict(s)
+        model.to(device)
+    if dtype:
+        model = model.to(dtype=dtype)
+    return model
+
+
+def normalized(a: torch.Tensor, order=2, dim=-1):
+    l2 = a.norm(order, dim, keepdim=True)
+    l2[l2 == 0] = 1
+    return a / l2
+
+
+@torch.no_grad()
+def encode_images(images: List[Image.Image], model2, preprocess, device='cuda') -> torch.Tensor:
+    if isinstance(images, Image.Image):
+        images = [images]
+    image_tensors = [preprocess(img).unsqueeze(0) for img in images]
+    image_batch = torch.cat(image_tensors).to(device)
+    image_features = model2.encode_image(image_batch)
+    im_emb_arr = normalized(image_features).cpu().float()
+    return im_emb_arr

modules/utils.py

@@ -0,0 +1,10 @@
+from huggingface_hub import hf_hub_download
+
+
+def download_from_url(url, cache_dir=None, verbose=True):
+    split = url.split("/")
+    username, repo_id, model_name = split[-3], split[-2], split[-1]
+    if verbose:
+        print(f"downloading: {username}/{repo_id}/{model_name}")
+    model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
+    return model_path

requirements.txt

@@ -0,0 +1,5 @@
+torch
+clip
+pytorch-lightning
+pillow
+huggingface-hub