Upload tagger.py

tagger.py

@@ -49,6 +49,18 @@ DANBOORU_TO_E621_RATING_MAP = {
 }
 
 
+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)
+
+
 def load_dict_from_csv(filename):
     with open(filename, 'r', encoding="utf-8") as f:
         lines = f.readlines()
@@ -227,20 +239,19 @@ def convert_tags_to_ja(input_prompt: str = ""):
     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)
+enable_auto_recom_prompt = True
 
-    animagine_ps = to_list("anime artwork, anime style, key visual, vibrant, studio anime, highly detailed, 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")
+
+animagine_ps = to_list("anime artwork, anime style, vibrant, studio anime, highly detailed, 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")
+other_ps = to_list("anime artwork, anime style, vibrant, studio anime, highly detailed, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed")
+other_nps = to_list("photo, deformed, black and white, realism, disfigured, low contrast, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly")
+default_ps = to_list("score_9, score_8_up, score_7_up, highly detailed, masterpiece, best quality, very aesthetic, absurdres")
+default_nps = to_list("score_6, score_5, score_4, 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]")
+def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
+    global enable_auto_recom_prompt
     prompts = to_list(prompt)
     neg_prompts = to_list(neg_prompt)
 
@@ -250,12 +261,16 @@ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "Non
     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
+    if type == "Auto":
+        enable_auto_recom_prompt = True
+    else:
+        enable_auto_recom_prompt = False
+        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)
@@ -263,6 +278,43 @@ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "Non
     return prompt, neg_prompt
 
 
+def load_model_prompt_dict():
+    import json
+    dict = {}
+    try:
+        with open('model_dict.json', encoding='utf-8') as f:
+            dict = json.load(f)
+    except Exception:
+        pass
+    return dict
+
+
+model_prompt_dict = load_model_prompt_dict()
+
+
+def insert_model_recom_prompt(prompt: str = "", neg_prompt: str = "", model_name: str = "None"):
+    if not model_name or not enable_auto_recom_prompt: return prompt, neg_prompt
+    prompts = to_list(prompt)
+    neg_prompts = to_list(neg_prompt)
+    prompts = list_sub(prompts, animagine_ps + pony_ps + other_ps)
+    neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps + other_nps)
+    last_empty_p = [""] if not prompts and type != "None" else []
+    last_empty_np = [""] if not neg_prompts and type != "None" else []
+    ps = []
+    nps = []
+    if model_name in model_prompt_dict.keys(): 
+        ps = to_list(model_prompt_dict[model_name]["prompt"])
+        nps = to_list(model_prompt_dict[model_name]["negative_prompt"])
+    else:
+        ps = default_ps
+        nps = default_nps
+    prompts = prompts + ps
+    neg_prompts = neg_prompts + 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')