Upload tagger.py

tagger.py

@@ -1,13 +1,13 @@
 from PIL import Image
 import torch
 import gradio as gr
-import spaces  # ZERO GPU
-
+import spaces
 from transformers import (
     AutoImageProcessor,
     AutoModelForImageClassification,
 )
 
+
 WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
 WD_MODEL_NAME = WD_MODEL_NAMES[0]
 
@@ -49,6 +49,34 @@ DANBOORU_TO_E621_RATING_MAP = {
 }
 
 
+# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+
+
+def replace_underline(x: str):
+    return x.strip().replace("_", " ") if x not in kaomojis else x.strip()
+
+
 def to_list(s):
     return [x.strip() for x in s.split(",") if not s == ""]
 
@@ -110,7 +138,7 @@ def select_random_character(series: str, character: str):
 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("_", " "), "")
+        etag = e621_dict.get(replace_underline(dtag), "")
         if etag:
             return etag
         else:
@@ -134,7 +162,7 @@ def convert_danbooru_to_e621_prompt(input_prompt: str = "", prompt_type: str = "
 
     e621_dict = danbooru_to_e621_dict
     for tag in tags:
-        tag = tag.strip().replace("_", " ")
+        tag = replace_underline(tag)
         tag = danbooru_to_e621(tag, e621_dict)
         if tag in PEOPLE_TAGS:        
             people_tags.append(tag)
@@ -162,6 +190,7 @@ def translate_prompt(prompt: str = ""):
             translated_prompt = translator.translate(prompt, src='auto', dest='en').text
             return translated_prompt
         except Exception as e:
+            print(e)
             return prompt
 
     def is_japanese(s):
@@ -194,6 +223,7 @@ def translate_prompt_to_ja(prompt: str = ""):
             translated_prompt = translator.translate(prompt, src='en', dest='ja').text
             return translated_prompt
         except Exception as e:
+            print(e)
             return prompt
 
     def is_japanese(s):
@@ -219,7 +249,7 @@ def translate_prompt_to_ja(prompt: str = ""):
 def tags_to_ja(itag, dict):
     def t_to_j(match, dict):
         tag = match.group(0)
-        ja = dict.get(tag.strip().replace("_", " "), "")
+        ja = dict.get(replace_underline(tag), "")
         if ja:
             return ja
         else:
@@ -238,7 +268,7 @@ def convert_tags_to_ja(input_prompt: str = ""):
     tags_to_ja_dict = load_dict_from_csv('all_tags_ja_ext.csv')
     dict = tags_to_ja_dict
     for tag in tags:
-        tag = tag.strip().replace("_", " ")
+        tag = replace_underline(tag)
         tag = tags_to_ja(tag, dict)
         out_tags.append(tag)
     
@@ -365,7 +395,7 @@ def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
 
     group_dict = tag_group_dict
     for tag in tags:
-        tag = tag.strip().replace("_", " ")
+        tag = replace_underline(tag)
         if tag in PEOPLE_TAGS:
             people_tags.append(tag)
         elif is_necessary(tag, keep_tags, group_dict):
@@ -393,7 +423,7 @@ def sort_taglist(tags: list[str]):
     rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
 
     for tag in tags:
-        tag = tag.strip().replace("_", " ")
+        tag = replace_underline(tag)
         if tag in PEOPLE_TAGS:
             people_tags.append(tag)
         elif tag in rating_set:
@@ -494,12 +524,13 @@ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_t
         output_series_tag = output_series_list[0]
     else:
         output_series_tag = ""
-    return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True),
+    return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True)
 
 
-def predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3, character_threshold: float = 0.8):
+def predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3,
+                     character_threshold: float = 0.8, input_series: str = "", input_character: str = ""):
     if not "Use WD Tagger" in algo and len(algo) != 0:
-        return "", "", input_tags, gr.update(interactive=True),
+        return input_series, input_character, input_tags, gr.update(interactive=True)
     return predict_tags(image, general_threshold, character_threshold)