update

.gitignore

@@ -0,0 +1,164 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+tmp.py

app.py

@@ -0,0 +1,28 @@
+import json
+import gradio as gr
+import os
+import spaces
+from tqdm import tqdm
+from PIL import Image
+from utils import WaifuScorer
+
+SCORER = WaifuScorer(
+    device='cuda',
+    verbose=True,
+)
+
+
+@spaces.GPU
+def score_image(image: Image.Image) -> float:
+    return SCORER([image])[0]
+
+
+demo = gr.Interface(
+    fn=score_image,
+    inputs=gr.Image(type='pil', label='Image'),
+    outputs=gr.Number(label='Score', precision=2),
+    title='Waifu Scorer V3',
+    description='''Score ranges from 0 to 10, higher is better.
+[Github](https://github.com/Eugeoter/waifu-scorer) | [Model](https://huggingface.co/Eugeoter/waifu-scorer-v3) | [Inspiration](https://github.com/christophschuhmann/improved-aesthetic-predictor)''',
+)
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+safetensors
+pytorch-lightning

utils.py

@@ -0,0 +1,128 @@
+import os
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+from typing import List, Union
+from PIL import Image
+
+WS_REPOS = ["Eugeoter/waifu-scorer-v3"]
+
+
+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)
+
+
+class WaifuScorer(object):
+    def __init__(self, model_path=None, device='cuda', cache_dir=None, verbose=False):
+        self.verbose = verbose
+        if model_path is None:
+            model_path = repo2path(WS_REPOS[0])
+            if self.verbose:
+                print(f"model path not set, switch to default: `{model_path}`")
+        if not os.path.isfile(model_path):
+            model_path = download_from_url(model_path, cache_dir=cache_dir)
+
+        print(f"loading pretrained model from `{model_path}`")
+        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()
+
+    @torch.no_grad()
+    def __call__(self, images: List[Image.Image]) -> Union[List[float], 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 repo2path(model_repo_and_path: str):
+    if os.path.isfile(model_repo_and_path):
+        model_path = model_repo_and_path
+    elif os.path.isdir(model_repo_and_path):
+        model_path = os.path.join(model_repo_and_path, "model.pth")
+    elif model_repo_and_path in WS_REPOS:
+        model_path = model_repo_and_path + '/model.pth'
+    else:
+        raise ValueError(f"Invalid model_repo_and_path: {model_repo_and_path}")
+    return model_path
+
+
+def download_from_url(url, cache_dir=None, verbose=True):
+    from huggingface_hub import hf_hub_download
+    split = url.split("/")
+    username, repo_id, model_name = split[-3], split[-2], split[-1]
+    # if verbose:
+    # print(f"[download_from_url]: {username}/{repo_id}/{model_name}")
+    model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
+    return model_path
+
+
+def load_clip_models(name: str = "ViT-L/14", device='cuda'):
+    import clip
+    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