diff --git a/userscripts/taggers/aesthetic_shadow.py b/userscripts/taggers/aesthetic_shadow.py index 96a1b53..4275fa4 100644 --- a/userscripts/taggers/aesthetic_shadow.py +++ b/userscripts/taggers/aesthetic_shadow.py @@ -15,6 +15,19 @@ from scripts.tagger import Tagger # I'm not sure if this is really working BATCH_SIZE = 3 +# tags used in Animagine-XL +SCORE_N = { + 'very aesthetic':0.71, + 'aesthetic':0.45, + 'displeasing':0.27, + 'very displeasing':-float('inf'), +} + +def get_aesthetic_tag(score:float): + for k, v in SCORE_N.items(): + if score > v: + return k + class AestheticShadowV2(Tagger): def load(self): if devices.device.index is None: @@ -40,8 +53,7 @@ class AestheticShadowV2(Tagger): for d in data: final[d["label"]] = d["score"] hq = final['hq'] - lq = final['lq'] - return [f"score_{math.floor((hq + (1 - lq))/2 * 10)}"] + return [get_aesthetic_tag(hq)] def predict(self, image: Image.Image, threshold=None): data = self.pipe_aesthetic(image) diff --git a/userscripts/taggers/cafeai_aesthetic_classifier.py b/userscripts/taggers/cafeai_aesthetic_classifier.py index 8464af6..99e96ea 100644 --- a/userscripts/taggers/cafeai_aesthetic_classifier.py +++ b/userscripts/taggers/cafeai_aesthetic_classifier.py @@ -36,9 +36,10 @@ class CafeAIAesthetic(Tagger): final = {} for d in data: final[d["label"]] = d["score"] - nae = final['not_aesthetic'] ae = final['aesthetic'] - return [f"score_{math.floor((ae + (1 - nae))/2 * 10)}"] + + # edit here to change tag + return [f"[CAFE]score_{math.floor(ae*10)}"] def predict(self, image: Image.Image, threshold=None): data = self.pipe_aesthetic(image, top_k=2) diff --git a/userscripts/taggers/improved_aesthetic_predictor.py b/userscripts/taggers/improved_aesthetic_predictor.py index fe217c5..eaa7199 100644 --- a/userscripts/taggers/improved_aesthetic_predictor.py +++ b/userscripts/taggers/improved_aesthetic_predictor.py @@ -69,7 +69,8 @@ class ImprovedAestheticPredictor(Tagger): def predict(self, image: Image.Image, threshold=None): image_embeds = image_embeddings(image, self.clip_model, self.clip_processor) prediction:torch.Tensor = self.model(torch.from_numpy(image_embeds).float().to(devices.device)) - return [f"score_{math.floor(prediction.item())}"] + # edit here to change tag + return [f"[IAP]score_{math.floor(prediction.item())}"] def name(self): return "Improved Aesthetic Predictor" \ No newline at end of file diff --git a/userscripts/taggers/waifu_aesthetic_classifier.py b/userscripts/taggers/waifu_aesthetic_classifier.py index 1952f51..5bbf66f 100644 --- a/userscripts/taggers/waifu_aesthetic_classifier.py +++ b/userscripts/taggers/waifu_aesthetic_classifier.py @@ -67,7 +67,8 @@ class WaifuAesthetic(Tagger): def predict(self, image: Image.Image, threshold=None): image_embeds = image_embeddings(image, self.clip_model, self.clip_processor) prediction:torch.Tensor = self.model(torch.from_numpy(image_embeds).float().to(devices.device)) - return [f"score_{math.floor(prediction.item()*10)}"] + # edit here to change tag + return [f"[WD]score_{math.floor(prediction.item()*10)}"] def name(self): return "wd aesthetic classifier" \ No newline at end of file