Aesthetic classifier training & inference tools.

This repo contains pretrained models that are to be used for rating anime-styled images with a 0 to 1 score, where the lowest score means the image has a low aesthetic rating and a high score means that the image has a high aesthetic rating.

⚙️ Pretrained Models

Name	CLIP Model	Performance
aes-B32-v0	OpenAI B-32	0.0266

🔑 Setup

git clone git@github.com:waifu-diffusion/aesthetic.git
cd aesthetic
python -m pip install -r requirements

🤖 Inference

import torch
from transformers import CLIPModel, CLIPProcessor
from aesthetic import image_embeddings, Classifier

aesthetic_path = 'aes-B32-v0.pth'
clip_name = 'openai/clip-vit-base-patch32'
url = 'https://cdn.donmai.us/original/16/67/__klein_moretti_lord_of_the_mysteries_drawn_by_ji26725339__1667415282975e8f8c574ca26d83e3be.jpg'

clipprocessor = CLIPProcessor.from_pretrained(clip_name)
clipmodel = CLIPModel.from_pretrained(clip_name).to('cuda').eval()

aes_model = Classifier(512, 256, 1)
aes_model.load_state_dict(torch.load(aesthetic_path))

image_embeds = image_embeddings(url, clipmodel, clipprocessor)
prediction = aes_model(torch.from_numpy(image_embeds).float().to('cuda'))
print(f'Prediction: {prediction.item()}')

And this should be the output:

$ python3 test.py 
Prediction: 0.999903678894043