John6666 commited on
Commit
80e6c51
0 Parent(s):

Super-squash branch 'main' using huggingface_hub

Browse files
.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Natural Text to SD Prompt Translator With LLM alpha
3
+ emoji: 👀😻
4
+ colorFrom: red
5
+ colorTo: purple
6
+ sdk: gradio
7
+ sdk_version: 4.38.1
8
+ app_file: app.py
9
+ pinned: false
10
+ license: apache-2.0
11
+ ---
12
+
13
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from utils import (
3
+ gradio_copy_text,
4
+ COPY_ACTION_JS,
5
+ )
6
+ from tagger import (
7
+ convert_danbooru_to_e621_prompt,
8
+ insert_recom_prompt,
9
+ )
10
+ from genimage import (
11
+ generate_image,
12
+ )
13
+ from llmdolphin import (
14
+ get_llm_formats,
15
+ get_dolphin_model_format,
16
+ get_dolphin_models,
17
+ get_dolphin_model_info,
18
+ select_dolphin_model,
19
+ select_dolphin_format,
20
+ add_dolphin_models,
21
+ get_dolphin_sysprompt,
22
+ get_dolphin_sysprompt_mode,
23
+ select_dolphin_sysprompt,
24
+ get_dolphin_languages,
25
+ select_dolphin_language,
26
+ dolphin_respond,
27
+ dolphin_parse,
28
+ )
29
+
30
+
31
+ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css="") as app:
32
+ gr.Markdown("""# Natural Text to SD Prompt Translator With LLM alpha
33
+ Text in natural language (English, Japanese, ...) => Prompt
34
+ """)
35
+ with gr.Column(scale=1):
36
+ with gr.Group():
37
+ chatbot = gr.Chatbot(likeable=False, show_copy_button=True, show_share_button=False, layout="panel", container=True, )
38
+ with gr.Row():
39
+ chat_msg = gr.Textbox(show_label=False, placeholder="Input text in English, Japanese, or any other languages and press Enter or click Send.", scale=4)
40
+ chat_submit = gr.Button("Send", scale=1)
41
+ chat_clear = gr.Button("Clear", scale=1)
42
+ with gr.Accordion("Additional inputs", open=False):
43
+ chat_format = gr.Dropdown(choices=get_llm_formats(), value=get_dolphin_model_format(get_dolphin_models()[0][1]), label="Message format")
44
+ chat_sysmsg = gr.Textbox(value=get_dolphin_sysprompt(), label="System message")
45
+ chat_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max tokens")
46
+ chat_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
47
+ chat_topp = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
48
+ chat_topk = gr.Slider(minimum=0, maximum=100, value=40, step=1, label="Top-k")
49
+ chat_rp = gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty")
50
+ with gr.Accordion("Add models", open=True):
51
+ chat_add_text = gr.Textbox(label="URL or Repo ID", placeholder="http://huggingface.co/.../...gguf or author/model", lines=1)
52
+ chat_add_format = gr.Dropdown(choices=get_llm_formats(), value=get_llm_formats()[0], label="Message format")
53
+ chat_add_submit = gr.Button("Update lists of models")
54
+ with gr.Accordion("Modes", open=True):
55
+ chat_model = gr.Dropdown(choices=get_dolphin_models(), value=get_dolphin_models()[0][1], allow_custom_value=True, label="Model")
56
+ chat_model_info = gr.Markdown(value=get_dolphin_model_info(get_dolphin_models()[0][1]), label="Model info")
57
+ with gr.Row():
58
+ chat_mode = gr.Dropdown(choices=get_dolphin_sysprompt_mode(), value=get_dolphin_sysprompt_mode()[0], allow_custom_value=False, label="Mode")
59
+ chat_lang = gr.Dropdown(choices=get_dolphin_languages(), value="English", allow_custom_value=True, label="Output language")
60
+
61
+ with gr.Column(scale=1):
62
+ with gr.Row():
63
+ with gr.Group():
64
+ output_text = gr.TextArea(label="Output tags", interactive=False, show_copy_button=True)
65
+ copy_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
66
+ elapsed_time_md = gr.Markdown(label="Elapsed time", value="", visible=False)
67
+ with gr.Group():
68
+ output_text_pony = gr.TextArea(label="Output tags (Pony e621 style)", interactive=False, show_copy_button=True)
69
+ copy_btn_pony = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
70
+ with gr.Accordion(label="Advanced options", open=False, visible=False):
71
+ tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False)
72
+ dummy_np = gr.Textbox(label="Negative prompt", value="", visible=False)
73
+ dummy_np_pony = gr.Textbox(label="Negative prompt", value="", visible=False)
74
+ recom_animagine = gr.Textbox(label="Animagine reccomended prompt", value="Animagine", visible=False)
75
+ recom_pony = gr.Textbox(label="Pony reccomended prompt", value="Pony", visible=False)
76
+ generate_image_btn = gr.Button(value="GENERATE IMAGE", size="lg", variant="primary")
77
+ result_image = gr.Gallery(label="Generated images", columns=1, object_fit="contain", container=True, preview=True, show_label=False, show_share_button=False, show_download_button=True, interactive=False, visible=True, format="png")
78
+
79
+ gr.on(
80
+ triggers=[chat_msg.submit, chat_submit.click],
81
+ fn=dolphin_respond,
82
+ inputs=[chat_msg, chatbot, chat_model, chat_sysmsg, chat_tokens, chat_temperature, chat_topp, chat_topk, chat_rp],
83
+ outputs=[chatbot],
84
+ queue=True,
85
+ show_progress="full",
86
+ trigger_mode="once",
87
+ ).success(dolphin_parse, [chatbot], [output_text, copy_btn, copy_btn_pony]).success(
88
+ convert_danbooru_to_e621_prompt, [output_text, tag_type], [output_text_pony], queue=False,
89
+ ).success(
90
+ insert_recom_prompt, [output_text, dummy_np, recom_animagine], [output_text, dummy_np], queue=False,
91
+ ).success(
92
+ insert_recom_prompt, [output_text_pony, dummy_np_pony, recom_pony], [output_text_pony, dummy_np_pony], queue=False,
93
+ )
94
+ chat_clear.click(lambda: None, None, chatbot, queue=False)
95
+ chat_model.change(select_dolphin_model, [chat_model], [chat_model, chat_format, chat_model_info], queue=True, show_progress="full")\
96
+ .success(lambda: None, None, chatbot, queue=False)
97
+ chat_format.change(select_dolphin_format, [chat_format], [chat_format], queue=False)\
98
+ .success(lambda: None, None, chatbot, queue=False)
99
+ chat_mode.change(select_dolphin_sysprompt, [chat_mode], [chat_sysmsg], queue=False)
100
+ chat_lang.change(select_dolphin_language, [chat_lang], [chat_sysmsg], queue=False)
101
+ gr.on(
102
+ triggers=[chat_add_text.submit, chat_add_submit.click],
103
+ fn=add_dolphin_models,
104
+ inputs=[chat_add_text, chat_add_format],
105
+ outputs=[chat_model],
106
+ queue=False,
107
+ trigger_mode="once",
108
+ )
109
+
110
+ copy_btn.click(gradio_copy_text, [output_text], js=COPY_ACTION_JS)
111
+ copy_btn_pony.click(gradio_copy_text, [output_text_pony], js=COPY_ACTION_JS)
112
+
113
+ generate_image_btn.click(generate_image, [output_text, dummy_np], [result_image], show_progress="full")
114
+
115
+
116
+ if __name__ == "__main__":
117
+ app.queue()
118
+ app.launch()
character_series_dict.csv ADDED
The diff for this file is too large to render. See raw diff
 
danbooru_e621.csv ADDED
The diff for this file is too large to render. See raw diff
 
genimage.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+
3
+
4
+ def load_pipeline():
5
+ from diffusers import DiffusionPipeline
6
+ import torch
7
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
8
+ pipe = DiffusionPipeline.from_pretrained(
9
+ "John6666/rae-diffusion-xl-v2-sdxl-spo-pcm",
10
+ custom_pipeline="lpw_stable_diffusion_xl",
11
+ torch_dtype=torch.float16,
12
+ )
13
+ pipe.to(device)
14
+ return pipe
15
+
16
+
17
+ def save_image(image, metadata, output_dir):
18
+ import os
19
+ import uuid
20
+ import json
21
+ from PIL import PngImagePlugin
22
+ filename = str(uuid.uuid4()) + ".png"
23
+ os.makedirs(output_dir, exist_ok=True)
24
+ filepath = os.path.join(output_dir, filename)
25
+ metadata_str = json.dumps(metadata)
26
+ info = PngImagePlugin.PngInfo()
27
+ info.add_text("metadata", metadata_str)
28
+ image.save(filepath, "PNG", pnginfo=info)
29
+ return filepath
30
+
31
+
32
+ pipe = load_pipeline()
33
+
34
+
35
+ @spaces.GPU
36
+ def generate_image(prompt, neg_prompt):
37
+ metadata = {
38
+ "prompt": prompt,
39
+ "negative_prompt": neg_prompt,
40
+ "resolution": f"{1024} x {1024}",
41
+ "guidance_scale": 7.5,
42
+ "num_inference_steps": 16,
43
+ "sampler": "Euler",
44
+ }
45
+ try:
46
+ images = pipe(
47
+ prompt=prompt,
48
+ prompt_2="anime artwork, anime style, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres",
49
+ negative_prompt=neg_prompt,
50
+ negative_prompt_2="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], photo, deformed, disfigured, low contrast, photo, deformed, disfigured, low contrast",
51
+ width=1024,
52
+ height=1024,
53
+ guidance_scale=7.5,
54
+ num_inference_steps=16,
55
+ output_type="pil",
56
+ clip_skip=1,
57
+ ).images
58
+ if images:
59
+ image_paths = [
60
+ save_image(image, metadata, "./outputs")
61
+ for image in images
62
+ ]
63
+ return image_paths
64
+ except Exception as e:
65
+ return []
66
+
ja_to_danbooru/character_series_dict.json ADDED
The diff for this file is too large to render. See raw diff
 
ja_to_danbooru/danbooru_tagtype_dict.json ADDED
The diff for this file is too large to render. See raw diff
 
ja_to_danbooru/ja_danbooru_dict.json ADDED
The diff for this file is too large to render. See raw diff
 
ja_to_danbooru/ja_to_danbooru.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import re
3
+ from pathlib import Path
4
+
5
+
6
+ def load_json_dict(path: str):
7
+ import json
8
+ from pathlib import Path
9
+ dict = {}
10
+ if not Path(path).exists(): return dict
11
+ try:
12
+ with open(path, encoding='utf-8') as f:
13
+ dict = json.load(f)
14
+ except Exception:
15
+ print(f"Failed to open dictionary file: {path}")
16
+ return dict
17
+ return dict
18
+
19
+
20
+ ja_danbooru_dict = load_json_dict('ja_danbooru_dict.json')
21
+ char_series_dict = load_json_dict('character_series_dict.json')
22
+ tagtype_dict = load_json_dict('danbooru_tagtype_dict.json')
23
+
24
+
25
+ def jatags_to_danbooru_tags(jatags: list[str]):
26
+ from rapidfuzz.process import extractOne
27
+ from rapidfuzz.utils import default_process
28
+ keys = list(ja_danbooru_dict.keys())
29
+ ckeys = list(char_series_dict.keys())
30
+ tags = []
31
+ for jatag in jatags:
32
+ jatag = str(jatag).strip()
33
+ s = default_process(str(jatag))
34
+ e1 = extractOne(s, keys, processor=default_process, score_cutoff=90.0)
35
+ if e1:
36
+ tag = str(ja_danbooru_dict[e1[0]])
37
+ tags.append(tag)
38
+ if tag in tagtype_dict.keys() and tagtype_dict[tag] == "character":
39
+ cs = default_process(tag)
40
+ ce1 = extractOne(cs, ckeys, processor=default_process, score_cutoff=95.0)
41
+ if ce1:
42
+ series = str(char_series_dict[ce1[0]])
43
+ tags.append(series)
44
+ return tags
45
+
46
+
47
+ def jatags_to_danbooru(input_tag, input_file, output_file, is_append):
48
+ if input_file and Path(input_file).exists():
49
+ try:
50
+ with open(input_file, 'r', encoding='utf-8') as f:
51
+ input_tag = f.read()
52
+ except Exception:
53
+ print(f"Failed to open input file: {input_file}")
54
+ ja_tags = [tag.strip() for tag in input_tag.split(",")] if input_tag else []
55
+ tags = jatags_to_danbooru_tags(ja_tags)
56
+ output_tags = ja_tags + tags if is_append else tags
57
+ output_tag = ", ".join(output_tags)
58
+ if output_file:
59
+ try:
60
+ with open(output_file, mode='w', encoding="utf-8") as f:
61
+ f.write(output_tag)
62
+ except Exception:
63
+ print(f"Failed to write output file: {output_file}")
64
+ else:
65
+ print(output_tag)
66
+ return output_tag
67
+
68
+
69
+ if __name__ == "__main__":
70
+ parser = argparse.ArgumentParser()
71
+ parser.add_argument("--tags", default=None, type=str, required=False, help="Input tags.")
72
+ parser.add_argument("--file", default=None, type=str, required=False, help="Input tags from a text file.")
73
+ parser.add_argument("--out", default=None, type=str, help="Output to text file.")
74
+ parser.add_argument("--append", default=False, type=bool, help="Whether the output contains the input tags or not.")
75
+
76
+ args = parser.parse_args()
77
+ assert (args.tags, args.file) != (None, None), "Must provide --tags or --file!"
78
+
79
+ jatags_to_danbooru(args.tags, args.file, args.out, args.append)
80
+
81
+
82
+ # Usage:
83
+ # python ja_to_danbooru.py --tags "女の子, 大室櫻子"
84
+ # python danbooru_to_ja.py --file inputtag.txt
85
+ # python danbooru_to_ja.py --file inputtag.txt --append True
86
+ # Datasets: https://huggingface.co/datasets/p1atdev/danbooru-ja-tag-pair-20240715
87
+ # Datasets: https://github.com/ponapon280/danbooru-e621-converter
llmdolphin.py ADDED
@@ -0,0 +1,527 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from llama_cpp import Llama
4
+ from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
5
+ from llama_cpp_agent.providers import LlamaCppPythonProvider
6
+ from llama_cpp_agent.chat_history import BasicChatHistory
7
+ from llama_cpp_agent.chat_history.messages import Roles
8
+ from ja_to_danbooru.ja_to_danbooru import jatags_to_danbooru_tags
9
+
10
+
11
+ llm_models_dir = "./llm_models"
12
+ llm_models = {
13
+ "L3-8B-Celeste-v1-Q5_K_M.gguf": ["bartowski/L3-8B-Celeste-v1-GGUF", MessagesFormatterType.LLAMA_3],
14
+ "L3-8B-Celeste-V1.2-Q5_K_M.gguf": ["bartowski/L3-8B-Celeste-V1.2-GGUF", MessagesFormatterType.LLAMA_3],
15
+ "Llama-3-Nymeria-ELYZA-8B.i1-Q4_K_M.gguf": ["mradermacher/Llama-3-Nymeria-ELYZA-8B-i1-GGUF", MessagesFormatterType.LLAMA_3],
16
+ "suzume-llama-3-8B-japanese.Q4_K_M.gguf": ["PrunaAI/lightblue-suzume-llama-3-8B-japanese-GGUF-smashed", MessagesFormatterType.LLAMA_3],
17
+ "suzume-llama-3-8B-multilingual-orpo-borda-top25.Q4_K_M.gguf": ["RichardErkhov/lightblue_-_suzume-llama-3-8B-multilingual-orpo-borda-top25-gguf", MessagesFormatterType.LLAMA_3],
18
+ "gemma-2-9b-it-SimPO.i1-Q4_K_M.gguf": ["mradermacher/gemma-2-9b-it-SimPO-i1-GGUF", MessagesFormatterType.ALPACA],
19
+ "Gemma-2-9B-It-SPPO-Iter3.Q4_K_M.iMatrix.gguf": ["MCZK/Gemma-2-9B-It-SPPO-Iter3-GGUF", MessagesFormatterType.ALPACA],
20
+ "Llama-3-NeuralPaca-8b.Q4_K_M.gguf": ["RichardErkhov/NeuralNovel_-_Llama-3-NeuralPaca-8b-gguf", MessagesFormatterType.ALPACA],
21
+ "SaoRPM-2x8B.i1-Q4_K_M.gguf": ["mradermacher/SaoRPM-2x8B-i1-GGUF", MessagesFormatterType.LLAMA_3],
22
+ "L3-Hecate-8B-v1.2.Q4_K_M.gguf": ["mradermacher/L3-Hecate-8B-v1.2-GGUF", MessagesFormatterType.LLAMA_3],
23
+ "Mahou-1.3b-llama3-8B.i1-Q4_K_M.gguf": ["mradermacher/Mahou-1.3b-llama3-8B-i1-GGUF", MessagesFormatterType.LLAMA_3],
24
+ "SwallowMaid-8B-L3-SPPO-abliterated.i1-Q5_K_M.gguf": ["mradermacher/SwallowMaid-8B-L3-SPPO-abliterated-i1-GGUF", MessagesFormatterType.LLAMA_3],
25
+ "L3-8B-Lunar-Stheno.i1-Q5_K_M.gguf": ["mradermacher/L3-8B-Lunar-Stheno-i1-GGUF", MessagesFormatterType.LLAMA_3],
26
+ "llama3_Loradent.Q4_K_M.gguf": ["mradermacher/llama3_Loradent-GGUF", MessagesFormatterType.LLAMA_3],
27
+ "Llama-3-8B-Stroganoff.i1-Q4_K_M.gguf": ["mradermacher/Llama-3-8B-Stroganoff-i1-GGUF", MessagesFormatterType.LLAMA_3],
28
+ "L3-8B-EnchantedForest-v0.5.i1-Q4_K_M.gguf": ["mradermacher/L3-8B-EnchantedForest-v0.5-i1-GGUF", MessagesFormatterType.LLAMA_3],
29
+ "gemma-radiation-rp-9b-q5_k_m.gguf": ["pegasus912/Gemma-Radiation-RP-9B-Q5_K_M-GGUF", MessagesFormatterType.MISTRAL],
30
+ "Magic-Dolphin-7b.Q4_K_M.gguf": ["mradermacher/Magic-Dolphin-7b-GGUF", MessagesFormatterType.MISTRAL],
31
+ "mathstral-7B-v0.1-Q5_K_M.gguf": ["bartowski/mathstral-7B-v0.1-GGUF", MessagesFormatterType.MISTRAL],
32
+ "Gemma2-9B-it-Boku-v1.Q5_K_M.gguf": ["mradermacher/Gemma2-9B-it-Boku-v1-GGUF", MessagesFormatterType.MISTRAL],
33
+ "Gemma-2-9B-It-SPPO-Iter3-Q5_K_M.gguf": ["grapevine-AI/Gemma-2-9B-It-SPPO-Iter3-GGUF", MessagesFormatterType.MISTRAL],
34
+ "L3-8B-Niitama-v1.i1-Q4_K_M.gguf": ["mradermacher/L3-8B-Niitama-v1-i1-GGUF", MessagesFormatterType.LLAMA_3],
35
+ "Maidphin-Kunoichi-7B.Q5_K_M.gguf": ["RichardErkhov/nbeerbower_-_Maidphin-Kunoichi-7B-gguf", MessagesFormatterType.MISTRAL],
36
+ "L3-15B-EtherealMaid-t0.0001.i1-Q4_K_M.gguf": ["mradermacher/L3-15B-EtherealMaid-t0.0001-i1-GGUF", MessagesFormatterType.LLAMA_3],
37
+ "L3-15B-MythicalMaid-t0.0001.i1-Q4_K_M.gguf": ["mradermacher/L3-15B-MythicalMaid-t0.0001-i1-GGUF", MessagesFormatterType.LLAMA_3],
38
+ "llama-3-Nephilim-v3-8B.Q5_K_M.gguf": ["grimjim/llama-3-Nephilim-v3-8B-GGUF", MessagesFormatterType.LLAMA_3],
39
+ "NarutoDolphin-10B.Q5_K_M.gguf": ["RichardErkhov/FelixChao_-_NarutoDolphin-10B-gguf", MessagesFormatterType.MISTRAL],
40
+ "l3-8b-tamamo-v1-q8_0.gguf": ["Ransss/L3-8B-Tamamo-v1-Q8_0-GGUF", MessagesFormatterType.LLAMA_3],
41
+ "Tiger-Gemma-9B-v1-Q4_K_M.gguf": ["bartowski/Tiger-Gemma-9B-v1-GGUF", MessagesFormatterType.LLAMA_3],
42
+ "TooManyMixRolePlay-7B-Story_V3.5.Q4_K_M.gguf": ["mradermacher/TooManyMixRolePlay-7B-Story_V3.5-GGUF", MessagesFormatterType.LLAMA_3],
43
+ "natsumura-llama3-v1.1-8b.Q4_K_M.gguf": ["mradermacher/natsumura-llama3-v1.1-8b-GGUF", MessagesFormatterType.LLAMA_3],
44
+ "natsumura-llama3-v1-8b.i1-Q4_K_M.gguf": ["mradermacher/natsumura-llama3-v1-8b-i1-GGUF", MessagesFormatterType.LLAMA_3],
45
+ "nephra_v1.0.Q5_K_M.gguf": ["PrunaAI/yodayo-ai-nephra_v1.0-GGUF-smashed", MessagesFormatterType.LLAMA_3],
46
+ "DPO-ONLY-Zephyr-7B.Q6_K.gguf": ["mradermacher/DPO-ONLY-Zephyr-7B-GGUF", MessagesFormatterType.LLAMA_3],
47
+ "L3-Deluxe-Scrambled-Eggs-On-Toast-8B.Q8_0.gguf": ["mradermacher/L3-Deluxe-Scrambled-Eggs-On-Toast-8B-GGUF", MessagesFormatterType.LLAMA_3],
48
+ "L3-Scrambled-Eggs-On-Toast-8B.i1-Q6_K.gguf": ["mradermacher/L3-Scrambled-Eggs-On-Toast-8B-i1-GGUF", MessagesFormatterType.LLAMA_3],
49
+ "Llama-3-uncensored-Dare-1.Q4_K_M.gguf": ["mradermacher/Llama-3-uncensored-Dare-1-GGUF", MessagesFormatterType.LLAMA_3],
50
+ "llama3-8B-DarkIdol-2.2-Uncensored-1048K.i1-Q6_K.gguf": ["mradermacher/llama3-8B-DarkIdol-2.2-Uncensored-1048K-i1-GGUF", MessagesFormatterType.LLAMA_3],
51
+ "dolphin-2.9.3-mistral-7b-32k-q4_k_m.gguf": ["huggingkot/dolphin-2.9.3-mistral-7B-32k-Q4_K_M-GGUF", MessagesFormatterType.MISTRAL],
52
+ "dolphin-2.9.3-mistral-7B-32k-Q5_K_M.gguf": ["bartowski/dolphin-2.9.3-mistral-7B-32k-GGUF", MessagesFormatterType.MISTRAL],
53
+ "Lexi-Llama-3-8B-Uncensored_Q5_K_M.gguf": ["Orenguteng/Llama-3-8B-Lexi-Uncensored-GGUF", MessagesFormatterType.LLAMA_3],
54
+ "Llama3-Sophie.Q8_0.gguf": ["mradermacher/Llama3-Sophie-GGUF", MessagesFormatterType.LLAMA_3],
55
+ "Aura-Uncensored-OAS-8B-L3.i1-Q4_K_M.gguf": ["mradermacher/Aura-Uncensored-OAS-8B-L3-i1-GGUF", MessagesFormatterType.LLAMA_3],
56
+ "L3-Uncen-Merger-Omelette-RP-v0.2-8B-Q5_K_S-imat.gguf": ["LWDCLS/L3-Uncen-Merger-Omelette-RP-v0.2-8B-GGUF-IQ-Imatrix-Request", MessagesFormatterType.LLAMA_3],
57
+ "qwen2-diffusion-prompter-v01-q6_k.gguf": ["trollek/Qwen2-0.5B-DiffusionPrompter-v0.1-GGUF", MessagesFormatterType.LLAMA_3],
58
+ "Smegmma-Deluxe-9B-v1-Q6_K.gguf": ["bartowski/Smegmma-Deluxe-9B-v1-GGUF", MessagesFormatterType.MISTRAL],
59
+ "Mahou-1.3c-mistral-7B.i1-Q6_K.gguf": ["mradermacher/Mahou-1.3c-mistral-7B-i1-GGUF", MessagesFormatterType.MISTRAL],
60
+ "Silicon-Maid-7B-Q8_0_X.gguf": ["duyntnet/Silicon-Maid-7B-imatrix-GGUF", MessagesFormatterType.ALPACA],
61
+ "l3-umbral-mind-rp-v3.0-8b-q5_k_m-imat.gguf": ["Casual-Autopsy/L3-Umbral-Mind-RP-v3.0-8B-Q5_K_M-GGUF", MessagesFormatterType.LLAMA_3],
62
+ "Phi-3.1-mini-128k-instruct-Q6_K_L.gguf": ["bartowski/Phi-3.1-mini-128k-instruct-GGUF", MessagesFormatterType.PHI_3],
63
+ "tifa-7b-qwen2-v0.1.q4_k_m.gguf": ["Tifa-RP/Tifa-7B-Qwen2-v0.1-GGUF", MessagesFormatterType.OPEN_CHAT],
64
+ "Llama-3-EZO-8b-Common-it.Q5_K_M.iMatrix.gguf": ["MCZK/Llama-3-EZO-8b-Common-it-GGUF", MessagesFormatterType.MISTRAL],
65
+ "EZO-Common-9B-gemma-2-it.i1-Q4_K_M.gguf": ["mradermacher/EZO-Common-9B-gemma-2-it-i1-GGUF", MessagesFormatterType.MISTRAL],
66
+ #"": ["", MessagesFormatterType.LLAMA_3],
67
+ #"": ["", MessagesFormatterType.MISTRAL],
68
+ #"": ["", MessagesFormatterType.ALPACA],
69
+ #"": ["", MessagesFormatterType.OPEN_CHAT],
70
+ }
71
+ llm_formats = {
72
+ "MISTRAL": MessagesFormatterType.MISTRAL,
73
+ "CHATML": MessagesFormatterType.CHATML,
74
+ "VICUNA": MessagesFormatterType.VICUNA,
75
+ "LLAMA 2": MessagesFormatterType.LLAMA_2,
76
+ "SYNTHIA": MessagesFormatterType.SYNTHIA,
77
+ "NEURAL CHAT": MessagesFormatterType.NEURAL_CHAT,
78
+ "SOLAR": MessagesFormatterType.SOLAR,
79
+ "OPEN CHAT": MessagesFormatterType.OPEN_CHAT,
80
+ "ALPACA": MessagesFormatterType.ALPACA,
81
+ "CODE DS": MessagesFormatterType.CODE_DS,
82
+ "B22": MessagesFormatterType.B22,
83
+ "LLAMA 3": MessagesFormatterType.LLAMA_3,
84
+ "PHI 3": MessagesFormatterType.PHI_3,
85
+ "Autocoder": MessagesFormatterType.AUTOCODER,
86
+ "DeepSeek Coder v2": MessagesFormatterType.DEEP_SEEK_CODER_2,
87
+ "Gemma 2": MessagesFormatterType.ALPACA,
88
+ "Qwen2": MessagesFormatterType.OPEN_CHAT,
89
+ }
90
+ # https://github.com/Maximilian-Winter/llama-cpp-agent
91
+ llm_languages = ["English", "Japanese", "Chinese"]
92
+ llm_models_tupled_list = []
93
+ default_llm_model_filename = list(llm_models.keys())[0]
94
+ override_llm_format = None
95
+
96
+
97
+ def to_list(s):
98
+ return [x.strip() for x in s.split(",") if not s == ""]
99
+
100
+
101
+ def list_uniq(l):
102
+ return sorted(set(l), key=l.index)
103
+
104
+
105
+ def to_list_ja(s):
106
+ import re
107
+ s = re.sub(r'[、。]', ',', s)
108
+ return [x.strip() for x in s.split(",") if not s == ""]
109
+
110
+
111
+ def is_japanese(s):
112
+ import unicodedata
113
+ for ch in s:
114
+ name = unicodedata.name(ch, "")
115
+ if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
116
+ return True
117
+ return False
118
+
119
+
120
+ def update_llm_model_tupled_list():
121
+ from pathlib import Path
122
+ global llm_models_tupled_list
123
+ llm_models_tupled_list = []
124
+ for k, v in llm_models.items():
125
+ name = k
126
+ value = k
127
+ llm_models_tupled_list.append((name, value))
128
+ model_files = Path(llm_models_dir).glob('*.gguf')
129
+ for path in model_files:
130
+ name = path.name
131
+ value = path.name
132
+ llm_models_tupled_list.append((name, value))
133
+ llm_models_tupled_list = list_uniq(llm_models_tupled_list)
134
+ return llm_models_tupled_list
135
+
136
+
137
+ def download_llm_models():
138
+ from huggingface_hub import hf_hub_download
139
+ global llm_models_tupled_list
140
+ llm_models_tupled_list = []
141
+ for k, v in llm_models.items():
142
+ try:
143
+ hf_hub_download(repo_id = v[0], filename = k, local_dir = llm_models_dir)
144
+ except Exception:
145
+ continue
146
+ name = k
147
+ value = k
148
+ llm_models_tupled_list.append((name, value))
149
+
150
+
151
+ def download_llm_model(filename):
152
+ from huggingface_hub import hf_hub_download
153
+ if not filename in llm_models.keys(): return default_llm_model_filename
154
+ try:
155
+ hf_hub_download(repo_id = llm_models[filename][0], filename = filename, local_dir = llm_models_dir)
156
+ except Exception:
157
+ return default_llm_model_filename
158
+ update_llm_model_tupled_list()
159
+ return filename
160
+
161
+
162
+ def get_dolphin_model_info(filename):
163
+ md = "None"
164
+ items = llm_models.get(filename, None)
165
+ if items:
166
+ md = f'Repo: [{items[0]}](https://huggingface.co/{items[0]})'
167
+ return md
168
+
169
+
170
+ def select_dolphin_model(filename, progress=gr.Progress(track_tqdm=True)):
171
+ global override_llm_format
172
+ override_llm_format = None
173
+ progress(0, desc="Loading model...")
174
+ value = download_llm_model(filename)
175
+ progress(1, desc="Model loaded.")
176
+ md = get_dolphin_model_info(filename)
177
+ return gr.update(value=value, choices=get_dolphin_models()), gr.update(value=get_dolphin_model_format(value)), gr.update(value=md)
178
+
179
+
180
+ def select_dolphin_format(format_name):
181
+ global override_llm_format
182
+ override_llm_format = llm_formats[format_name]
183
+ return gr.update(value=format_name)
184
+
185
+
186
+ #download_llm_models()
187
+ download_llm_model(default_llm_model_filename)
188
+
189
+
190
+ def get_dolphin_models():
191
+ return update_llm_model_tupled_list()
192
+
193
+
194
+ def get_llm_formats():
195
+ return list(llm_formats.keys())
196
+
197
+
198
+ def get_key_from_value(d, val):
199
+ keys = [k for k, v in d.items() if v == val]
200
+ if keys:
201
+ return keys[0]
202
+ return None
203
+
204
+
205
+ def get_dolphin_model_format(filename):
206
+ if not filename in llm_models.keys(): filename = default_llm_model_filename
207
+ format = llm_models[filename][1]
208
+ format_name = get_key_from_value(llm_formats, format)
209
+ return format_name
210
+
211
+
212
+ def add_dolphin_models(query, format_name):
213
+ import re
214
+ from huggingface_hub import HfApi
215
+ global llm_models
216
+ api = HfApi()
217
+ add_models = {}
218
+ format = llm_formats[format_name]
219
+ filename = ""
220
+ repo = ""
221
+ try:
222
+ s = list(re.findall(r'^(?:https?://huggingface.co/)?(.+?/.+?)(?:/.*/(.+?.gguf).*?)?$', query)[0])
223
+ if s and "" in s: s.remove("")
224
+ if len(s) == 1:
225
+ repo = s[0]
226
+ if not api.repo_exists(repo_id = repo): return gr.update(visible=True)
227
+ files = api.list_repo_files(repo_id = repo)
228
+ for file in files:
229
+ if str(file).endswith(".gguf"): add_models[filename] = [repo, format]
230
+ elif len(s) >= 2:
231
+ repo = s[0]
232
+ filename = s[1]
233
+ if not api.repo_exists(repo_id = repo) or not api.file_exists(repo_id = repo, filename = filename): return gr.update(visible=True)
234
+ add_models[filename] = [repo, format]
235
+ else: return gr.update(visible=True)
236
+ except Exception:
237
+ return gr.update(visible=True)
238
+ print(add_models)
239
+ llm_models = (llm_models | add_models).copy()
240
+ return gr.update(choices=get_dolphin_models())
241
+
242
+
243
+ dolphin_output_language = "English"
244
+ dolphin_sysprompt_mode = "Default"
245
+ dolphin_system_prompt = {"Default": r'''You are a helpful AI assistant to generate messages for AI that outputs an image when I enter a message.
246
+ The message must have the following [Tags] generated in strict accordance with the following [Rules]:
247
+ ```
248
+ [Tags]
249
+ - Words to describe full names of characters and names of series in which they appear.
250
+ - Words to describe names of the people there and their numbers, such as 2girls, 1boy.
251
+ - Words to describe their hair color, hairstyle, hair length, hair accessory, eye color, eye shape, facial expression, breast size, and clothing of them in detail, such as long hair.
252
+ - Words to describe their external features, ornaments and belongings (also specify colors, patterns, shapes) in detail.
253
+ - Words to describe their stance from head to toe in detail.
254
+ - Words to describe their acting, especially with sexual activity in detail.
255
+ - Words to describe their surroundings in detail.
256
+ - Words to describe background details, such as inside room, forest, starry sky.
257
+ [Rules]
258
+ - Any output should be plain text in English and don't use line breaks.
259
+ - Output only composed of Tags in 1 line, separated by commas with spaces between Tags, in lower case English.
260
+ - Output should be in the format: "//GENBEGIN//://1girl, Tag, Tag, ..., Tag//://GENEND//".
261
+ - Preferably refer to and describe the information obtained from Danbooru. If not, describe it in own way.
262
+ - It's preferable that each Tag is a plain phrase, word, caption, Danbooru tag, or E621 tag.
263
+ - Convert any nicknames to full names first.
264
+ - If a sexual theme is given, priority should be given to specific and rich descriptions of sexual activity, especially about genitals, fluids.
265
+ - Assemble a short story internally which is developed from the themes provided, then describe a scene into an detailed English sentences based on the central character internally.
266
+ - Split sentences into short phrases or words, and then convert them to Tags.
267
+ - Use associated Danbooru tags, E621 tags.
268
+ - Same Tags should be used only once per output.
269
+ - Anyway, keep processing until you've finished outputting message.
270
+ ```
271
+ Based on these Rules, please tell me message within 40 Tags that can generate an image for the following themes:
272
+ ''',
273
+ "With dialogue and description": r'''You are a helpful AI assistant to generate messages for AI that outputs an image when I enter a message.
274
+ The message must have the following [Tags] generated in strict accordance with the following [Rules]:
275
+ ```
276
+ [Tags]
277
+ - Words to describe full names of characters and names of series in which they appear.
278
+ - Words to describe names of the people there and their numbers, such as 2girls, 1boy.
279
+ - Words to describe their hair color, hairstyle, hair length, hair accessory, eye color, eye shape, facial expression, breast size, and clothing of them in detail, such as long hair.
280
+ - Words to describe their external features, ornaments and belongings (also specify colors, patterns, shapes) in detail.
281
+ - Words to describe their stance from head to toe in detail.
282
+ - Words to describe their acting, especially with sexual activity in detail.
283
+ - Words to describe their surroundings in detail.
284
+ - Words to describe background details, such as inside room, forest, starry sky.
285
+ [Rules]
286
+ - Any Tags should be plain text in English and don't use line breaks.
287
+ - Message is only composed of Tags in 1 line, separated by commas with spaces between Tags, in lower case English.
288
+ - Message should be in the format: "//GENBEGIN//://1girl, Tag, Tag, ..., Tag//://GENEND//".
289
+ - Preferably refer to and describe the information obtained from Danbooru. If not, describe it in own way.
290
+ - It's preferable that each Tag is a plain phrase, word, caption, Danbooru tag, or E621 tag.
291
+ - Convert any nicknames to full names first.
292
+ - If a sexual theme is given, priority should be given to specific and rich descriptions of sexual activity, especially about genitals, fluids.
293
+ - Assemble a short story internally which is developed from the themes provided, then describe a scene into an detailed English sentences based on the central character internally.
294
+ - Split sentences into short phrases or words, and then convert them to Tags.
295
+ - Use associated Danbooru tags, E621 tags.
296
+ - Same Tags should be used only once per output.
297
+ - Anyway, keep processing until you've finished outputting message.
298
+ ```
299
+ Based on these Rules, please tell me message within 40 Tags that can generate an image for the following themes,
300
+ then write the character's long actor's line composed of one's voices and moaning and voices in thought, based on the story you have assembled, in <LANGUAGE> only,
301
+ enclosed in //VOICEBEGIN//:// and //://VOICEEND//, then describe the message you've generated in short, in <LANGUAGE> only.:
302
+ ''', "Japanese to Danbooru Dictionary": r"""You are a helpful AI assistant.
303
+ Extract Japanese words from the following sentences and output them separated by commas. Convert words in their original forms.
304
+ Output should be enclosed in //GENBEGIN//:// and //://GENEND//. The text to be given is as follows:""",
305
+ "Chat with LLM": r"You are a helpful AI assistant. Respond in <LANGUAGE>."}
306
+
307
+
308
+ def get_dolphin_sysprompt():
309
+ import re
310
+ prompt = re.sub('<LANGUAGE>', dolphin_output_language, dolphin_system_prompt.get(dolphin_sysprompt_mode, ""))
311
+ return prompt
312
+
313
+
314
+ def get_dolphin_sysprompt_mode():
315
+ return list(dolphin_system_prompt.keys())
316
+
317
+
318
+ def select_dolphin_sysprompt(key: str):
319
+ global dolphin_sysprompt_mode
320
+ if not key in dolphin_system_prompt.keys():
321
+ dolphin_sysprompt_mode = "Default"
322
+ else:
323
+ dolphin_sysprompt_mode = key
324
+ return gr.update(value=get_dolphin_sysprompt())
325
+
326
+
327
+ def get_dolphin_languages():
328
+ return llm_languages
329
+
330
+
331
+ def select_dolphin_language(lang: str):
332
+ global dolphin_output_language
333
+ dolphin_output_language = lang
334
+ return gr.update(value=get_dolphin_sysprompt())
335
+
336
+
337
+ @spaces.GPU
338
+ def dolphin_respond(
339
+ message: str,
340
+ history: list[tuple[str, str]],
341
+ model: str = default_llm_model_filename,
342
+ system_message: str = get_dolphin_sysprompt(),
343
+ max_tokens: int = 1024,
344
+ temperature: float = 0.7,
345
+ top_p: float = 0.95,
346
+ top_k: int = 40,
347
+ repeat_penalty: float = 1.1,
348
+ progress=gr.Progress(track_tqdm=True),
349
+ ):
350
+ from pathlib import Path
351
+ progress(0, desc="Processing...")
352
+
353
+ if override_llm_format:
354
+ chat_template = override_llm_format
355
+ else:
356
+ chat_template = llm_models[model][1]
357
+
358
+ llm = Llama(
359
+ model_path=str(Path(f"{llm_models_dir}/{model}")),
360
+ flash_attn=True,
361
+ n_gpu_layers=81,
362
+ n_batch=1024,
363
+ n_ctx=8192,
364
+ )
365
+ provider = LlamaCppPythonProvider(llm)
366
+
367
+ agent = LlamaCppAgent(
368
+ provider,
369
+ system_prompt=f"{system_message}",
370
+ predefined_messages_formatter_type=chat_template,
371
+ debug_output=False
372
+ )
373
+
374
+ settings = provider.get_provider_default_settings()
375
+ settings.temperature = temperature
376
+ settings.top_k = top_k
377
+ settings.top_p = top_p
378
+ settings.max_tokens = max_tokens
379
+ settings.repeat_penalty = repeat_penalty
380
+ settings.stream = True
381
+
382
+ messages = BasicChatHistory()
383
+
384
+ for msn in history:
385
+ user = {
386
+ 'role': Roles.user,
387
+ 'content': msn[0]
388
+ }
389
+ assistant = {
390
+ 'role': Roles.assistant,
391
+ 'content': msn[1]
392
+ }
393
+ messages.add_message(user)
394
+ messages.add_message(assistant)
395
+
396
+ stream = agent.get_chat_response(
397
+ message,
398
+ llm_sampling_settings=settings,
399
+ chat_history=messages,
400
+ returns_streaming_generator=True,
401
+ print_output=False
402
+ )
403
+
404
+ progress(0.5, desc="Processing...")
405
+
406
+ outputs = ""
407
+ for output in stream:
408
+ outputs += output
409
+ yield [(outputs, None)]
410
+
411
+
412
+ def dolphin_parse(
413
+ history: list[tuple[str, str]],
414
+ ):
415
+ import re
416
+ if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: "", gr.update(visible=True), gr.update(visible=True)
417
+ try:
418
+ msg = history[-1][0]
419
+ except Exception:
420
+ return ""
421
+ m = re.findall(r'/GENBEGIN/((?:.|\s)+?)/GENEND/', msg)
422
+ raw_prompt = re.sub(r'[*/:_"#]|\n', ' ', ", ".join(m)).lower() if m else ""
423
+ prompts = []
424
+ if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
425
+ prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])
426
+ else:
427
+ prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
428
+ return ", ".join(prompts), gr.update(interactive=True), gr.update(interactive=True)
429
+
430
+
431
+ @spaces.GPU
432
+ def dolphin_respond_auto(
433
+ message: str,
434
+ history: list[tuple[str, str]],
435
+ model: str = default_llm_model_filename,
436
+ system_message: str = get_dolphin_sysprompt(),
437
+ max_tokens: int = 1024,
438
+ temperature: float = 0.7,
439
+ top_p: float = 0.95,
440
+ top_k: int = 40,
441
+ repeat_penalty: float = 1.1,
442
+ progress=gr.Progress(track_tqdm=True),
443
+ ):
444
+ #if not is_japanese(message): return [(None, None)]
445
+
446
+ from pathlib import Path
447
+ progress(0, desc="Processing...")
448
+
449
+ if override_llm_format:
450
+ chat_template = override_llm_format
451
+ else:
452
+ chat_template = llm_models[model][1]
453
+
454
+ llm = Llama(
455
+ model_path=str(Path(f"{llm_models_dir}/{model}")),
456
+ flash_attn=True,
457
+ n_gpu_layers=81,
458
+ n_batch=1024,
459
+ n_ctx=8192,
460
+ )
461
+ provider = LlamaCppPythonProvider(llm)
462
+
463
+ agent = LlamaCppAgent(
464
+ provider,
465
+ system_prompt=f"{system_message}",
466
+ predefined_messages_formatter_type=chat_template,
467
+ debug_output=False
468
+ )
469
+
470
+ settings = provider.get_provider_default_settings()
471
+ settings.temperature = temperature
472
+ settings.top_k = top_k
473
+ settings.top_p = top_p
474
+ settings.max_tokens = max_tokens
475
+ settings.repeat_penalty = repeat_penalty
476
+ settings.stream = True
477
+
478
+ messages = BasicChatHistory()
479
+
480
+ for msn in history:
481
+ user = {
482
+ 'role': Roles.user,
483
+ 'content': msn[0]
484
+ }
485
+ assistant = {
486
+ 'role': Roles.assistant,
487
+ 'content': msn[1]
488
+ }
489
+ messages.add_message(user)
490
+ messages.add_message(assistant)
491
+
492
+ progress(0, desc="Translating...")
493
+ stream = agent.get_chat_response(
494
+ message,
495
+ llm_sampling_settings=settings,
496
+ chat_history=messages,
497
+ returns_streaming_generator=True,
498
+ print_output=False
499
+ )
500
+
501
+ progress(0.5, desc="Processing...")
502
+
503
+ outputs = ""
504
+ for output in stream:
505
+ outputs += output
506
+ yield [(outputs, None)]
507
+
508
+
509
+ def dolphin_parse_simple(
510
+ message: str,
511
+ history: list[tuple[str, str]],
512
+ ):
513
+ import re
514
+ #if not is_japanese(message) or not history or len(history) < 1: return message
515
+ if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: return message
516
+ try:
517
+ msg = history[-1][0]
518
+ except Exception:
519
+ return ""
520
+ m = re.findall(r'/GENBEGIN/((?:.|\s)+?)/GENEND/', msg)
521
+ raw_prompt = re.sub(r'[*/:_"#]|\n', ' ', ", ".join(m)).lower() if m else ""
522
+ prompts = []
523
+ if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
524
+ prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])
525
+ else:
526
+ prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
527
+ return ", ".join(prompts)
pre-requirements.txt ADDED
@@ -0,0 +1 @@
 
 
1
+ pip>=23.0.0
requirements.txt ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ huggingface_hub
2
+ scikit-build-core
3
+ https://github.com/abetlen/llama-cpp-python/releases/download/v0.2.82-cu124/llama_cpp_python-0.2.82-cp310-cp310-linux_x86_64.whl
4
+ git+https://github.com/Maximilian-Winter/llama-cpp-agent
5
+ pybind11>=2.12
6
+ torch
7
+ torchvision
8
+ accelerate
9
+ transformers
10
+ optimum[onnxruntime]
11
+ spaces
12
+ dartrs
13
+ httpx==0.13.3
14
+ httpcore
15
+ googletrans==4.0.0rc1
16
+ git+https://github.com/huggingface/diffusers
17
+ rapidfuzz
tag_group.csv ADDED
The diff for this file is too large to render. See raw diff
 
tagger.py ADDED
@@ -0,0 +1,506 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from PIL import Image
2
+ import torch
3
+ import gradio as gr
4
+ import spaces # ZERO GPU
5
+
6
+ from transformers import (
7
+ AutoImageProcessor,
8
+ AutoModelForImageClassification,
9
+ )
10
+
11
+ WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
12
+ WD_MODEL_NAME = WD_MODEL_NAMES[0]
13
+
14
+ wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
15
+ wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
16
+ wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
17
+
18
+
19
+ def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
20
+ return (
21
+ [f"1{noun}"]
22
+ + [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
23
+ + [f"{maximum+1}+{noun}s"]
24
+ )
25
+
26
+
27
+ PEOPLE_TAGS = (
28
+ _people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
29
+ )
30
+
31
+
32
+ RATING_MAP = {
33
+ "general": "safe",
34
+ "sensitive": "sensitive",
35
+ "questionable": "nsfw",
36
+ "explicit": "explicit, nsfw",
37
+ }
38
+ DANBOORU_TO_E621_RATING_MAP = {
39
+ "safe": "rating_safe",
40
+ "sensitive": "rating_safe",
41
+ "nsfw": "rating_explicit",
42
+ "explicit, nsfw": "rating_explicit",
43
+ "explicit": "rating_explicit",
44
+ "rating:safe": "rating_safe",
45
+ "rating:general": "rating_safe",
46
+ "rating:sensitive": "rating_safe",
47
+ "rating:questionable, nsfw": "rating_explicit",
48
+ "rating:explicit, nsfw": "rating_explicit",
49
+ }
50
+
51
+
52
+ def to_list(s):
53
+ return [x.strip() for x in s.split(",") if not s == ""]
54
+
55
+
56
+ def list_sub(a, b):
57
+ return [e for e in a if e not in b]
58
+
59
+
60
+ def list_uniq(l):
61
+ return sorted(set(l), key=l.index)
62
+
63
+
64
+ def load_dict_from_csv(filename):
65
+ with open(filename, 'r', encoding="utf-8") as f:
66
+ lines = f.readlines()
67
+ dict = {}
68
+ for line in lines:
69
+ parts = line.strip().split(',')
70
+ dict[parts[0]] = parts[1]
71
+ return dict
72
+
73
+
74
+ anime_series_dict = load_dict_from_csv('character_series_dict.csv')
75
+
76
+
77
+ def character_list_to_series_list(character_list):
78
+ output_series_tag = []
79
+ series_tag = ""
80
+ series_dict = anime_series_dict
81
+ for tag in character_list:
82
+ series_tag = series_dict.get(tag, "")
83
+ if tag.endswith(")"):
84
+ tags = tag.split("(")
85
+ character_tag = "(".join(tags[:-1])
86
+ if character_tag.endswith(" "):
87
+ character_tag = character_tag[:-1]
88
+ series_tag = tags[-1].replace(")", "")
89
+
90
+ if series_tag:
91
+ output_series_tag.append(series_tag)
92
+
93
+ return output_series_tag
94
+
95
+
96
+ def select_random_character(series: str, character: str):
97
+ from random import randrange
98
+ character_list = list(anime_series_dict.keys())
99
+ character = character_list[randrange(len(character_list) - 1)]
100
+ series = anime_series_dict.get(character.split(",")[0].strip(), "")
101
+ return series, character
102
+
103
+
104
+ def danbooru_to_e621(dtag, e621_dict):
105
+ def d_to_e(match, e621_dict):
106
+ dtag = match.group(0)
107
+ etag = e621_dict.get(dtag.strip().replace("_", " "), "")
108
+ if etag:
109
+ return etag
110
+ else:
111
+ return dtag
112
+
113
+ import re
114
+ tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2)
115
+
116
+ return tag
117
+
118
+
119
+ danbooru_to_e621_dict = load_dict_from_csv('danbooru_e621.csv')
120
+
121
+
122
+ def convert_danbooru_to_e621_prompt(input_prompt: str = "", prompt_type: str = "danbooru"):
123
+ if prompt_type == "danbooru": return input_prompt
124
+ tags = input_prompt.split(",") if input_prompt else []
125
+ people_tags: list[str] = []
126
+ other_tags: list[str] = []
127
+ rating_tags: list[str] = []
128
+
129
+ e621_dict = danbooru_to_e621_dict
130
+ for tag in tags:
131
+ tag = tag.strip().replace("_", " ")
132
+ tag = danbooru_to_e621(tag, e621_dict)
133
+ if tag in PEOPLE_TAGS:
134
+ people_tags.append(tag)
135
+ elif tag in DANBOORU_TO_E621_RATING_MAP.keys():
136
+ rating_tags.append(DANBOORU_TO_E621_RATING_MAP.get(tag.replace(" ",""), ""))
137
+ else:
138
+ other_tags.append(tag)
139
+
140
+ rating_tags = sorted(set(rating_tags), key=rating_tags.index)
141
+ rating_tags = [rating_tags[0]] if rating_tags else []
142
+ rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
143
+
144
+ output_prompt = ", ".join(people_tags + other_tags + rating_tags)
145
+
146
+ return output_prompt
147
+
148
+
149
+ def translate_prompt(prompt: str = ""):
150
+ def translate_to_english(prompt):
151
+ import httpcore
152
+ setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
153
+ from googletrans import Translator
154
+ translator = Translator()
155
+ try:
156
+ translated_prompt = translator.translate(prompt, src='auto', dest='en').text
157
+ return translated_prompt
158
+ except Exception as e:
159
+ return prompt
160
+
161
+ def is_japanese(s):
162
+ import unicodedata
163
+ for ch in s:
164
+ name = unicodedata.name(ch, "")
165
+ if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
166
+ return True
167
+ return False
168
+
169
+ def to_list(s):
170
+ return [x.strip() for x in s.split(",")]
171
+
172
+ prompts = to_list(prompt)
173
+ outputs = []
174
+ for p in prompts:
175
+ p = translate_to_english(p) if is_japanese(p) else p
176
+ outputs.append(p)
177
+
178
+ return ", ".join(outputs)
179
+
180
+
181
+ def translate_prompt_to_ja(prompt: str = ""):
182
+ def translate_to_japanese(prompt):
183
+ import httpcore
184
+ setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
185
+ from googletrans import Translator
186
+ translator = Translator()
187
+ try:
188
+ translated_prompt = translator.translate(prompt, src='en', dest='ja').text
189
+ return translated_prompt
190
+ except Exception as e:
191
+ return prompt
192
+
193
+ def is_japanese(s):
194
+ import unicodedata
195
+ for ch in s:
196
+ name = unicodedata.name(ch, "")
197
+ if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
198
+ return True
199
+ return False
200
+
201
+ def to_list(s):
202
+ return [x.strip() for x in s.split(",")]
203
+
204
+ prompts = to_list(prompt)
205
+ outputs = []
206
+ for p in prompts:
207
+ p = translate_to_japanese(p) if not is_japanese(p) else p
208
+ outputs.append(p)
209
+
210
+ return ", ".join(outputs)
211
+
212
+
213
+ def tags_to_ja(itag, dict):
214
+ def t_to_j(match, dict):
215
+ tag = match.group(0)
216
+ ja = dict.get(tag.strip().replace("_", " "), "")
217
+ if ja:
218
+ return ja
219
+ else:
220
+ return tag
221
+
222
+ import re
223
+ tag = re.sub(r'[\w ]+', lambda wrapper: t_to_j(wrapper, dict), itag, 2)
224
+
225
+ return tag
226
+
227
+
228
+ def convert_tags_to_ja(input_prompt: str = ""):
229
+ tags = input_prompt.split(",") if input_prompt else []
230
+ out_tags = []
231
+
232
+ tags_to_ja_dict = load_dict_from_csv('all_tags_ja_ext.csv')
233
+ dict = tags_to_ja_dict
234
+ for tag in tags:
235
+ tag = tag.strip().replace("_", " ")
236
+ tag = tags_to_ja(tag, dict)
237
+ out_tags.append(tag)
238
+
239
+ return ", ".join(out_tags)
240
+
241
+
242
+ enable_auto_recom_prompt = True
243
+
244
+
245
+ animagine_ps = to_list("anime artwork, anime style, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres")
246
+ 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]")
247
+ pony_ps = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres")
248
+ 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")
249
+ other_ps = to_list("anime artwork, anime style, studio anime, highly detailed, cinematic photo, 35mm photograph, film, bokeh, professional, 4k, highly detailed")
250
+ other_nps = to_list("photo, deformed, black and white, realism, disfigured, low contrast, drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly")
251
+ default_ps = to_list("score_9, score_8_up, score_7_up, highly detailed, masterpiece, best quality, very aesthetic, absurdres")
252
+ 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]")
253
+ def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
254
+ global enable_auto_recom_prompt
255
+ prompts = to_list(prompt)
256
+ neg_prompts = to_list(neg_prompt)
257
+
258
+ prompts = list_sub(prompts, animagine_ps + pony_ps)
259
+ neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps)
260
+
261
+ last_empty_p = [""] if not prompts and type != "None" else []
262
+ last_empty_np = [""] if not neg_prompts and type != "None" else []
263
+
264
+ if type == "Auto":
265
+ enable_auto_recom_prompt = True
266
+ else:
267
+ enable_auto_recom_prompt = False
268
+ if type == "Animagine":
269
+ prompts = prompts + animagine_ps
270
+ neg_prompts = neg_prompts + animagine_nps
271
+ elif type == "Pony":
272
+ prompts = prompts + pony_ps
273
+ neg_prompts = neg_prompts + pony_nps
274
+
275
+ prompt = ", ".join(list_uniq(prompts) + last_empty_p)
276
+ neg_prompt = ", ".join(list_uniq(neg_prompts) + last_empty_np)
277
+
278
+ return prompt, neg_prompt
279
+
280
+
281
+ def load_model_prompt_dict():
282
+ import json
283
+ dict = {}
284
+ try:
285
+ with open('model_dict.json', encoding='utf-8') as f:
286
+ dict = json.load(f)
287
+ except Exception:
288
+ pass
289
+ return dict
290
+
291
+
292
+ model_prompt_dict = load_model_prompt_dict()
293
+
294
+
295
+ def insert_model_recom_prompt(prompt: str = "", neg_prompt: str = "", model_name: str = "None"):
296
+ if not model_name or not enable_auto_recom_prompt: return prompt, neg_prompt
297
+ prompts = to_list(prompt)
298
+ neg_prompts = to_list(neg_prompt)
299
+ prompts = list_sub(prompts, animagine_ps + pony_ps + other_ps)
300
+ neg_prompts = list_sub(neg_prompts, animagine_nps + pony_nps + other_nps)
301
+ last_empty_p = [""] if not prompts and type != "None" else []
302
+ last_empty_np = [""] if not neg_prompts and type != "None" else []
303
+ ps = []
304
+ nps = []
305
+ if model_name in model_prompt_dict.keys():
306
+ ps = to_list(model_prompt_dict[model_name]["prompt"])
307
+ nps = to_list(model_prompt_dict[model_name]["negative_prompt"])
308
+ else:
309
+ ps = default_ps
310
+ nps = default_nps
311
+ prompts = prompts + ps
312
+ neg_prompts = neg_prompts + nps
313
+ prompt = ", ".join(list_uniq(prompts) + last_empty_p)
314
+ neg_prompt = ", ".join(list_uniq(neg_prompts) + last_empty_np)
315
+ return prompt, neg_prompt
316
+
317
+
318
+ tag_group_dict = load_dict_from_csv('tag_group.csv')
319
+
320
+
321
+ def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
322
+ def is_dressed(tag):
323
+ import re
324
+ 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')
325
+ return p.search(tag)
326
+
327
+ def is_background(tag):
328
+ import re
329
+ p = re.compile(r'background|outline|light|sky|build|day|screen|tree|city')
330
+ return p.search(tag)
331
+
332
+ un_tags = ['solo']
333
+ group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
334
+ keep_group_dict = {
335
+ "body": ['groups', 'body_parts'],
336
+ "dress": ['groups', 'body_parts', 'attire'],
337
+ "all": group_list,
338
+ }
339
+
340
+ def is_necessary(tag, keep_tags, group_dict):
341
+ if keep_tags == "all":
342
+ return True
343
+ elif tag in un_tags or group_dict.get(tag, "") in explicit_group:
344
+ return False
345
+ elif keep_tags == "body" and is_dressed(tag):
346
+ return False
347
+ elif is_background(tag):
348
+ return False
349
+ else:
350
+ return True
351
+
352
+ if keep_tags == "all": return input_prompt
353
+ keep_group = keep_group_dict.get(keep_tags, keep_group_dict["body"])
354
+ explicit_group = list(set(group_list) ^ set(keep_group))
355
+
356
+ tags = input_prompt.split(",") if input_prompt else []
357
+ people_tags: list[str] = []
358
+ other_tags: list[str] = []
359
+
360
+ group_dict = tag_group_dict
361
+ for tag in tags:
362
+ tag = tag.strip().replace("_", " ")
363
+ if tag in PEOPLE_TAGS:
364
+ people_tags.append(tag)
365
+ elif is_necessary(tag, keep_tags, group_dict):
366
+ other_tags.append(tag)
367
+
368
+ output_prompt = ", ".join(people_tags + other_tags)
369
+
370
+ return output_prompt
371
+
372
+
373
+ def sort_taglist(tags: list[str]):
374
+ if not tags: return []
375
+ character_tags: list[str] = []
376
+ series_tags: list[str] = []
377
+ people_tags: list[str] = []
378
+ group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
379
+ group_tags = {}
380
+ other_tags: list[str] = []
381
+ rating_tags: list[str] = []
382
+
383
+ group_dict = tag_group_dict
384
+ group_set = set(group_dict.keys())
385
+ character_set = set(anime_series_dict.keys())
386
+ series_set = set(anime_series_dict.values())
387
+ rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
388
+
389
+ for tag in tags:
390
+ tag = tag.strip().replace("_", " ")
391
+ if tag in PEOPLE_TAGS:
392
+ people_tags.append(tag)
393
+ elif tag in rating_set:
394
+ rating_tags.append(tag)
395
+ elif tag in group_set:
396
+ elem = group_dict[tag]
397
+ group_tags[elem] = group_tags[elem] + [tag] if elem in group_tags else [tag]
398
+ elif tag in character_set:
399
+ character_tags.append(tag)
400
+ elif tag in series_set:
401
+ series_tags.append(tag)
402
+ else:
403
+ other_tags.append(tag)
404
+
405
+ output_group_tags: list[str] = []
406
+ for k in group_list:
407
+ output_group_tags.extend(group_tags.get(k, []))
408
+
409
+ rating_tags = [rating_tags[0]] if rating_tags else []
410
+ rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
411
+
412
+ output_tags = character_tags + series_tags + people_tags + output_group_tags + other_tags + rating_tags
413
+
414
+ return output_tags
415
+
416
+
417
+ def sort_tags(tags: str):
418
+ if not tags: return ""
419
+ taglist: list[str] = []
420
+ for tag in tags.split(","):
421
+ taglist.append(tag.strip())
422
+ taglist = list(filter(lambda x: x != "", taglist))
423
+ return ", ".join(sort_taglist(taglist))
424
+
425
+
426
+ def postprocess_results(results: dict[str, float], general_threshold: float, character_threshold: float):
427
+ results = {
428
+ k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
429
+ }
430
+
431
+ rating = {}
432
+ character = {}
433
+ general = {}
434
+
435
+ for k, v in results.items():
436
+ if k.startswith("rating:"):
437
+ rating[k.replace("rating:", "")] = v
438
+ continue
439
+ elif k.startswith("character:"):
440
+ character[k.replace("character:", "")] = v
441
+ continue
442
+
443
+ general[k] = v
444
+
445
+ character = {k: v for k, v in character.items() if v >= character_threshold}
446
+ general = {k: v for k, v in general.items() if v >= general_threshold}
447
+
448
+ return rating, character, general
449
+
450
+
451
+ def gen_prompt(rating: list[str], character: list[str], general: list[str]):
452
+ people_tags: list[str] = []
453
+ other_tags: list[str] = []
454
+ rating_tag = RATING_MAP[rating[0]]
455
+
456
+ for tag in general:
457
+ if tag in PEOPLE_TAGS:
458
+ people_tags.append(tag)
459
+ else:
460
+ other_tags.append(tag)
461
+
462
+ all_tags = people_tags + other_tags
463
+
464
+ return ", ".join(all_tags)
465
+
466
+
467
+ @spaces.GPU()
468
+ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
469
+ inputs = wd_processor.preprocess(image, return_tensors="pt")
470
+
471
+ outputs = wd_model(**inputs.to(wd_model.device, wd_model.dtype))
472
+ logits = torch.sigmoid(outputs.logits[0]) # take the first logits
473
+
474
+ # get probabilities
475
+ results = {
476
+ wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
477
+ }
478
+ # rating, character, general
479
+ rating, character, general = postprocess_results(
480
+ results, general_threshold, character_threshold
481
+ )
482
+ prompt = gen_prompt(
483
+ list(rating.keys()), list(character.keys()), list(general.keys())
<