Spaces:
Sleeping
Sleeping
from PIL import Image | |
import torch | |
from transformers import ( | |
AutoImageProcessor, | |
AutoModelForImageClassification, | |
) | |
import gradio as gr | |
import spaces # ZERO GPU | |
MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"] | |
MODEL_NAME = MODEL_NAMES[0] | |
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME, trust_remote_code=True) | |
model.to("cuda" if torch.cuda.is_available() else "cpu") | |
processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True) | |
def _people_tag(noun: str, minimum: int = 1, maximum: int = 5): | |
return ( | |
[f"1{noun}"] | |
+ [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)] | |
+ [f"{maximum+1}+{noun}s"] | |
) | |
PEOPLE_TAGS = ( | |
_people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"] | |
) | |
RATING_MAP = { | |
"general": "safe", | |
"sensitive": "sensitive", | |
"questionable": "nsfw", | |
"explicit": "explicit, nsfw", | |
} | |
DANBOORU_TO_E621_RATING_MAP = { | |
"safe": "rating_safe", | |
"sensitive": "rating_safe", | |
"nsfw": "rating_explicit", | |
"explicit, nsfw": "rating_explicit", | |
"explicit": "rating_explicit", | |
"rating:safe": "rating_safe", | |
"rating:general": "rating_safe", | |
"rating:sensitive": "rating_safe", | |
"rating:questionable, nsfw": "rating_explicit", | |
"rating:explicit, nsfw": "rating_explicit", | |
} | |
def load_dict_from_csv(filename): | |
with open(filename, 'r', encoding="utf-8") as f: | |
lines = f.readlines() | |
dict = {} | |
for line in lines: | |
parts = line.strip().split(',') | |
dict[parts[0]] = parts[1] | |
return dict | |
def get_series_dict(): | |
import re | |
with open('characterfull.txt', 'r') as f: | |
lines = f.readlines() | |
series_dict = {} | |
for line in lines: | |
parts = line.strip().split(', ') | |
if len(parts) >= 3: | |
name = parts[-2].replace("\\", "") | |
if name.endswith(")"): | |
names = name.split("(") | |
character_name = "(".join(names[:-1]) | |
if character_name.endswith(" "): | |
name = character_name[:-1] | |
series = re.sub(r'\\[()]', '', parts[-1]) | |
series_dict[name] = series | |
return series_dict | |
anime_series_dict = get_series_dict() | |
def character_list_to_series_list(character_list): | |
output_series_tag = [] | |
series_tag = "" | |
series_dict = anime_series_dict | |
for tag in character_list: | |
series_tag = series_dict.get(tag, "") | |
if tag.endswith(")"): | |
tags = tag.split("(") | |
character_tag = "(".join(tags[:-1]) | |
if character_tag.endswith(" "): | |
character_tag = character_tag[:-1] | |
series_tag = tags[-1].replace(")", "") | |
if series_tag: | |
output_series_tag.append(series_tag) | |
return output_series_tag | |
def danbooru_to_e621(dtag, e621_dict): | |
def d_to_e(match, e621_dict): | |
dtag = match.group(0) | |
etag = e621_dict.get(dtag.strip().replace("_", " "), "") | |
if etag: | |
return etag | |
else: | |
return dtag | |
import re | |
tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2) | |
return tag | |
danbooru_to_e621_dict = load_dict_from_csv('danbooru_e621.csv') | |
def convert_danbooru_to_e621_prompt(input_prompt: str = "", prompt_type: str = "danbooru"): | |
if prompt_type == "danbooru": return input_prompt | |
tags = input_prompt.split(",") if input_prompt else [] | |
people_tags: list[str] = [] | |
other_tags: list[str] = [] | |
rating_tags: list[str] = [] | |
e621_dict = danbooru_to_e621_dict | |
for tag in tags: | |
tag = tag.strip().replace("_", " ") | |
tag = danbooru_to_e621(tag, e621_dict) | |
if tag in PEOPLE_TAGS: | |
people_tags.append(tag) | |
elif tag in DANBOORU_TO_E621_RATING_MAP.keys(): | |
rating_tags.append(DANBOORU_TO_E621_RATING_MAP.get(tag.replace(" ",""), "")) | |
else: | |
other_tags.append(tag) | |
rating_tags = sorted(set(rating_tags), key=rating_tags.index) | |
rating_tags = [rating_tags[0]] if rating_tags else [] | |
rating_tags = ["explicit, nsfw"] if rating_tags[0] == "explicit" else rating_tags | |
output_prompt = ", ".join(people_tags + other_tags + rating_tags) | |
return output_prompt | |
def translate_prompt(prompt: str = ""): | |
def translate_to_english(prompt): | |
import httpcore | |
setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy') | |
from googletrans import Translator | |
translator = Translator() | |
try: | |
translated_prompt = translator.translate(prompt, src='auto', dest='en').text | |
return translated_prompt | |
except Exception as e: | |
return prompt | |
def is_japanese(s): | |
import unicodedata | |
for ch in s: | |
name = unicodedata.name(ch, "") | |
if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name: | |
return True | |
return False | |
def to_list(s): | |
return [x.strip() for x in s.split(",")] | |
prompts = to_list(prompt) | |
outputs = [] | |
for p in prompts: | |
p = translate_to_english(p) if is_japanese(p) else p | |
outputs.append(p) | |
return ", ".join(outputs) | |
def translate_prompt_to_ja(prompt: str = ""): | |
def translate_to_japanese(prompt): | |
import httpcore | |
setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy') | |
from googletrans import Translator | |
translator = Translator() | |
try: | |
translated_prompt = translator.translate(prompt, src='en', dest='ja').text | |
return translated_prompt | |
except Exception as e: | |
return prompt | |
def is_japanese(s): | |
import unicodedata | |
for ch in s: | |
name = unicodedata.name(ch, "") | |
if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name: | |
return True | |
return False | |
def to_list(s): | |
return [x.strip() for x in s.split(",")] | |
prompts = to_list(prompt) | |
outputs = [] | |
for p in prompts: | |
p = translate_to_japanese(p) if not is_japanese(p) else p | |
outputs.append(p) | |
return ", ".join(outputs) | |
def tags_to_ja(itag, dict): | |
def t_to_j(match, dict): | |
tag = match.group(0) | |
ja = dict.get(tag.strip().replace("_", " "), "") | |
if ja: | |
return ja | |
else: | |
return tag | |
import re | |
tag = re.sub(r'[\w ]+', lambda wrapper: t_to_j(wrapper, dict), itag, 2) | |
return tag | |
tags_to_ja_dict = load_dict_from_csv('all_tags_ja_ext.csv') | |
def convert_tags_to_ja(input_prompt: str = ""): | |
tags = input_prompt.split(",") if input_prompt else [] | |
out_tags = [] | |
dict = tags_to_ja_dict | |
for tag in tags: | |
tag = tag.strip().replace("_", " ") | |
tag = tags_to_ja(tag, dict) | |
out_tags.append(tag) | |
return ", ".join(out_tags) | |
def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"): | |
def to_list(s): | |
return [x.strip() for x in s.split(",") if not s == ""] | |
def list_sub(a, b): | |
return [e for e in a if e not in b] | |
def list_uniq(l): | |
return sorted(set(l), key=l.index) | |
animagine_ps = to_list("masterpiece, best quality, very aesthetic, absurdres") | |
animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]") | |
pony_ps = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres") | |
pony_nps = to_list("source_pony, source_furry, source_cartoon, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends") | |
prompts = to_list(prompt) | |
neg_prompts = to_list(neg_prompt) | |
prompts = list_sub(prompts, animagine_ps + pony_ps) | |
neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps) | |
last_empty_p = [""] if not prompts and type != "None" else [] | |
last_empty_np = [""] if not neg_prompts and type != "None" else [] | |
if type == "Animagine": | |
prompts = prompts + animagine_ps | |
neg_prompts = neg_prompts + animagine_nps | |
elif type == "Pony": | |
prompts = prompts + pony_ps | |
neg_prompts = neg_prompts + pony_nps | |
prompt = ", ".join(list_uniq(prompts) + last_empty_p) | |
neg_prompt = ", ".join(list_uniq(neg_prompts) + last_empty_np) | |
return prompt, neg_prompt | |
tag_group_dict = load_dict_from_csv('tag_group.csv') | |
def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"): | |
def is_dressed(tag): | |
import re | |
p = re.compile(r'dress|cloth|uniform|costume|vest|sweater|coat|shirt|jacket|blazer|apron|leotard|hood|sleeve|skirt|shorts|pant|loafer|ribbon|necktie|bow|collar|glove|sock|shoe|boots|wear|emblem') | |
return p.search(tag) | |
def is_background(tag): | |
import re | |
p = re.compile(r'background|outline|light|sky|build|day|screen|tree|city') | |
return p.search(tag) | |
un_tags = ['solo'] | |
group_list = ['people', 'age', 'pattern', 'place', 'hair', 'modifier', 'screen', 'animal', 'effect', 'situation', 'status', 'lighting', 'accesory', 'body', 'nsfw', 'camera', 'option', 'taste', 'other', 'detail', 'action', 'dress', 'character', 'face', 'costume', 'attribute', 'weather', 'temporary', 'gender', 'favorite', 'food', 'object', 'quality', 'expression', 'life', 'background'] | |
keep_group_dict = { | |
"body": ['people', 'age', 'hair', 'body', 'character', 'face', 'gender'], | |
"dress": ['people', 'age', 'hair', 'accesory', 'body', 'dress', 'character', 'face', 'costume', 'gender'], | |
"all": ['people', 'age', 'pattern', 'place', 'hair', 'modifier', 'screen', 'animal', 'effect', 'situation', 'status', 'lighting', 'accesory', 'body', 'nsfw', 'camera', 'option', 'taste', 'other', 'detail', 'action', 'dress', 'character', 'face', 'costume', 'attribute', 'weather', 'temporary', 'gender', 'favorite', 'food', 'object', 'quality', 'expression', 'life', 'background'] | |
} | |
def is_necessary(tag, keep_tags, group_dict): | |
if keep_tags == "all": | |
return True | |
elif tag in un_tags or group_dict.get(tag, "") in explicit_group: | |
return False | |
elif keep_tags == "body" and is_dressed(tag): | |
return False | |
elif is_background(tag): | |
return False | |
else: | |
return True | |
if keep_tags == "all": return input_prompt | |
keep_group = keep_group_dict.get(keep_tags, ['people', 'age', 'hair', 'body', 'character', 'face', 'gender']) | |
explicit_group = list(set(group_list) ^ set(keep_group)) | |
tags = input_prompt.split(",") if input_prompt else [] | |
people_tags: list[str] = [] | |
other_tags: list[str] = [] | |
group_dict = tag_group_dict | |
for tag in tags: | |
tag = tag.strip().replace("_", " ") | |
if tag in PEOPLE_TAGS: | |
people_tags.append(tag) | |
elif is_necessary(tag, keep_tags, group_dict): | |
other_tags.append(tag) | |
output_prompt = ", ".join(people_tags + other_tags) | |
return output_prompt | |
def postprocess_results(results: dict[str, float], general_threshold: float, character_threshold: float): | |
results = { | |
k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True) | |
} | |
rating = {} | |
character = {} | |
general = {} | |
for k, v in results.items(): | |
if k.startswith("rating:"): | |
rating[k.replace("rating:", "")] = v | |
continue | |
elif k.startswith("character:"): | |
character[k.replace("character:", "")] = v | |
continue | |
general[k] = v | |
character = {k: v for k, v in character.items() if v >= character_threshold} | |
general = {k: v for k, v in general.items() if v >= general_threshold} | |
return rating, character, general | |
def gen_prompt(rating: list[str], character: list[str], general: list[str]): | |
people_tags: list[str] = [] | |
other_tags: list[str] = [] | |
rating_tag = RATING_MAP[rating[0]] | |
for tag in general: | |
if tag in PEOPLE_TAGS: | |
people_tags.append(tag) | |
else: | |
other_tags.append(tag) | |
all_tags = people_tags + other_tags | |
return ", ".join(all_tags) | |
def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8): | |
inputs = processor.preprocess(image, return_tensors="pt") | |
outputs = model(**inputs.to(model.device, model.dtype)) | |
logits = torch.sigmoid(outputs.logits[0]) # take the first logits | |
# get probabilities | |
results = { | |
model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits) | |
} | |
# rating, character, general | |
rating, character, general = postprocess_results( | |
results, general_threshold, character_threshold | |
) | |
prompt = gen_prompt( | |
list(rating.keys()), list(character.keys()), list(general.keys()) | |
) | |
output_series_tag = "" | |
output_series_list = character_list_to_series_list(character.keys()) | |
if output_series_list: | |
output_series_tag = output_series_list[0] | |
else: | |
output_series_tag = "" | |
return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True), | |
def compose_prompt_to_copy(character: str, series: str, general: str): | |
characters = character.split(",") if character else [] | |
serieses = series.split(",") if series else [] | |
generals = general.split(",") if general else [] | |
tags = characters + serieses + generals | |
cprompt = ",".join(tags) if tags else "" | |
return cprompt | |