Initial version

README.md

@@ -1,11 +1,12 @@
 ---
 title: CityAesthetics Demo
-emoji: 🐨
+emoji: 🏙️
 colorFrom: blue
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.3.0
+sdk_version: 3.23.0
 app_file: app.py
+models: [city96/CityAesthetics]
 pinned: false
 license: apache-2.0
 ---

app.py

@@ -0,0 +1,113 @@
+import os
+import torch
+import gradio as gr
+from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from model import AestheticPredictorModel
+
+HFREPO = "City96/CityAesthetics"
+MODELS = [
+	"CityAesthetics-Anime-v1.8",
+]
+
+class CityAestheticsPipeline:
+	"""
+	Demo pipeline for [image=>score] prediction
+		Accepts a list of model paths on initialization.
+		Resulting object can be called directly with a PIL image as the input.
+		Returns a dict with the model name as key and the score [0.0;1.0] as a value.
+	"""
+	def __init__(self, model_paths):
+		self.models = {}
+		for path in model_paths:
+			name = os.path.splitext(os.path.basename(path))[0]
+			self.models[name] = self.load_model(path)
+
+		clip_ver = "openai/clip-vit-large-patch14"
+		self.proc = CLIPImageProcessor.from_pretrained(clip_ver)
+		self.clip = CLIPVisionModelWithProjection.from_pretrained(clip_ver)
+		print("CityAesthetics: Pipeline init ok") # debug
+
+	def load_model(self, path):
+		sd = load_file(path)
+		assert tuple(sd["up.0.weight"].shape) == (1024, 768) # only allow CLIP ver
+		model = AestheticPredictorModel()
+		model.load_state_dict(sd)
+		model.eval()
+		return model
+
+	def __call__(self, raw):
+		img = self.proc(images=raw, return_tensors="pt")
+		with torch.no_grad():
+			emb = self.clip(pixel_values=img["pixel_values"])
+		emb = emb["image_embeds"].detach().cpu()
+		out = {}
+		for name, model in self.models.items():
+			pred = model(emb)
+			out[name] = float(pred.squeeze(0))
+		return out
+
+def get_model_path(name):
+	fname = f"{name}.safetensors"
+
+	# local path: [models/AesPred-Anime-v1.8.safetensors]
+	path  = os.path.join(os.path.dirname(os.path.realpath(__file__)),"models")
+	if os.path.isfile(os.path.join(path, fname)):
+		print("CityAesthetics: Using local model")
+		return os.path.join(path, fname)
+
+	# huggingface hub fallback
+	print("CityAesthetics: Using HF Hub model")
+	return str(hf_hub_download(
+		token    = os.environ.get("HFS_TOKEN") or True,
+		repo_id  = HFREPO,
+		filename = fname,
+		# subfolder = fname.split('-')[1],
+	))
+
+article = """\
+# About
+
+This is the live demo for the CityAesthetics class of predictors.
+
+For more information, you can check out the [Huggingface Hub](https://huggingface.co/city96/CityAesthetics) or [GitHub page](https://github.com/city96/CityAesthetics).
+
+## CityAesthetics-Anime
+
+This flavor is optimized for scoring anime images with at least one subject present.
+
+### Intentional biases:
+
+- Completely negative towards real life photos (ideal score of 0%)
+- Strongly Negative towards text (subtitles, memes, etc) and manga panels
+- Fairly negative towards 3D and to some extent 2.5D images
+- Negative towards western cartoons and stylized images (chibi, parody)
+
+### Expected output scores:
+
+- Non-anime images should always score below 20%
+- Sketches/rough lineart/oekaki get around 20-40%
+- Flat shading/TV anime gets around 40-50%
+- Above 50% is mostly scored based on my personal style preferences
+
+### Issues:
+
+- Tends to filter male characters.
+- Requires at least 1 subject, won't work for scenery/landscapes.
+- Noticeable positive bias towards anime characters with animal ears.
+- Hit-or-miss with AI generated images due to style/quality not being correlated.
+"""
+
+pipeline = CityAestheticsPipeline([get_model_path(x) for x in MODELS])
+gr.Interface(
+	fn      = pipeline,
+	title   = "CityAesthetics demo",
+	article = article,
+	inputs  = gr.Image(label="Input image", type="pil"),
+	outputs = gr.Label(label="Model prediction", show_label=False),
+	examples = "./examples",
+	allow_flagging = "never",
+	analytics_enabled = False,
+).launch()

model.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+class ResBlock(nn.Module):
+	"""Block with residuals"""
+	def __init__(self, ch):
+		super().__init__()
+		self.join = nn.ReLU()
+		self.long = nn.Sequential(
+			nn.Linear(ch, ch),
+			nn.LeakyReLU(0.1),
+			nn.Linear(ch, ch),
+			nn.LeakyReLU(0.1),
+			nn.Linear(ch, ch),
+		)
+	def forward(self, x):
+		return self.join(self.long(x) + x)
+
+class AestheticPredictorModel(nn.Module):
+	"""
+	Main predictor class. Original:
+	https://github.com/city96/CityAesthetics/blob/main/model.py
+	"""
+	def __init__(self, features=768, hidden=1024):
+		super().__init__()
+		self.features = features
+		self.hidden = hidden
+		self.up = nn.Sequential(
+			nn.Linear(self.features, self.hidden),
+			ResBlock(ch=self.hidden),
+		)
+		self.down = nn.Sequential(
+			nn.Linear(self.hidden, 128),
+			nn.Linear(128, 64),
+			nn.Dropout(0.1),
+			nn.LeakyReLU(),
+			nn.Linear(64, 32),
+			nn.Linear(32,  1),
+			nn.Tanh(),
+		)
+	def forward(self, x):
+		y = self.up(x)
+		z = self.down(y)
+		return (z+1.0)/2.0

requirements.txt

@@ -0,0 +1,3 @@
+torch=2.1.0
+safetensors=0.4.0
+transformers=4.35.0