Super-squash branch 'main' using huggingface_hub

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Natural Text to SD Prompt Translator With LLM alpha
+emoji: 👀😻
+colorFrom: red
+colorTo: purple
+sdk: gradio
+sdk_version: 4.38.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,115 @@
+import gradio as gr
+from utils import (
+    gradio_copy_text,
+    COPY_ACTION_JS,
+)
+from tagger import (
+    convert_danbooru_to_e621_prompt,
+    insert_recom_prompt,
+)
+from genimage import (
+    generate_image,
+)
+from llmdolphin import (
+    get_llm_formats,
+    get_dolphin_model_format,
+    get_dolphin_models,
+    get_dolphin_model_info,
+    select_dolphin_model,
+    select_dolphin_format,
+    add_dolphin_models,
+    get_dolphin_sysprompt,
+    get_dolphin_sysprompt_mode,
+    select_dolphin_sysprompt,
+    get_dolphin_languages,
+    select_dolphin_language,
+    dolphin_respond,
+    dolphin_parse,
+)
+
+
+with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css="") as app:
+    gr.Markdown("""# Natural Text to SD Prompt Translator With LLM alpha
+                    Text in natural language (English, Japanese, ...) => Prompt
+                """)
+    with gr.Column(scale=1):
+        with gr.Group():
+            chatbot = gr.Chatbot(likeable=False, show_copy_button=True, show_share_button=False, layout="panel", container=True, )
+            with gr.Row():
+                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)
+                chat_submit = gr.Button("Send", scale=1)
+                chat_clear = gr.Button("Clear", scale=1)
+            with gr.Accordion("Additional inputs", open=False):
+                chat_model = gr.Dropdown(choices=get_dolphin_models(), value=get_dolphin_models()[0][1], allow_custom_value=True, label="Model")
+                chat_model_info = gr.Markdown(value=get_dolphin_model_info(get_dolphin_models()[0][1]), label="Model info")
+                chat_format = gr.Dropdown(choices=get_llm_formats(), value=get_dolphin_model_format(get_dolphin_models()[0][1]), label="Message format")
+                chat_sysmsg = gr.Textbox(value=get_dolphin_sysprompt(), label="System message")
+                chat_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max tokens")
+                chat_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+                chat_topp = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
+                chat_topk = gr.Slider(minimum=0, maximum=100, value=40, step=1, label="Top-k")
+                chat_rp = gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty")
+                with gr.Row():
+                    chat_mode = gr.Dropdown(choices=get_dolphin_sysprompt_mode(), value=get_dolphin_sysprompt_mode()[0], allow_custom_value=False, label="Mode")
+                    chat_lang = gr.Dropdown(choices=get_dolphin_languages(), value="English", allow_custom_value=True, label="Output language")
+                with gr.Accordion("Add models", open=False):
+                    chat_add_text = gr.Textbox(label="URL or Repo ID", placeholder="http://huggingface.co/.../...gguf or author/model", lines=1)
+                    chat_add_format = gr.Dropdown(choices=get_llm_formats(), value=get_llm_formats()[0], label="Message format")
+                    chat_add_submit = gr.Button("Update lists of models")
+
+    with gr.Column(scale=1):
+        with gr.Row():
+            with gr.Group():
+                output_text = gr.TextArea(label="Output tags", interactive=False, show_copy_button=True)
+                copy_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                elapsed_time_md = gr.Markdown(label="Elapsed time", value="", visible=False)
+            with gr.Group():
+                output_text_pony = gr.TextArea(label="Output tags (Pony e621 style)", interactive=False, show_copy_button=True)
+                copy_btn_pony = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+            with gr.Accordion(label="Advanced options", open=False, visible=False):
+                tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False)
+                dummy_np = gr.Textbox(label="Negative prompt", value="", visible=False)
+                dummy_np_pony = gr.Textbox(label="Negative prompt", value="", visible=False)
+                recom_animagine = gr.Textbox(label="Animagine reccomended prompt", value="Animagine", visible=False)
+                recom_pony = gr.Textbox(label="Pony reccomended prompt", value="Pony", visible=False)
+        generate_image_btn = gr.Button(value="GENERATE IMAGE", size="lg", variant="primary")
+        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")
+
+        gr.on(
+            triggers=[chat_msg.submit, chat_submit.click],
+            fn=dolphin_respond,
+            inputs=[chat_msg, chatbot, chat_model, chat_sysmsg, chat_tokens, chat_temperature, chat_topp, chat_topk, chat_rp],
+            outputs=[chatbot],
+            queue=True,
+            show_progress="full",
+            trigger_mode="once",
+        ).success(dolphin_parse, [chatbot], [output_text, copy_btn, copy_btn_pony]).success(
+            convert_danbooru_to_e621_prompt, [output_text, tag_type], [output_text_pony], queue=False,
+        ).success(
+            insert_recom_prompt, [output_text, dummy_np, recom_animagine], [output_text, dummy_np], queue=False,
+        ).success(
+            insert_recom_prompt, [output_text_pony, dummy_np_pony, recom_pony], [output_text_pony, dummy_np_pony], queue=False,
+        )
+        chat_clear.click(lambda: None, None, chatbot, queue=False)
+        chat_model.change(select_dolphin_model, [chat_model], [chat_model, chat_format, chat_model_info], queue=True, show_progress="full")
+        chat_format.change(select_dolphin_format, [chat_format], [chat_format], queue=False)
+        chat_mode.change(select_dolphin_sysprompt, [chat_mode], [chat_sysmsg], queue=False)
+        chat_lang.change(select_dolphin_language, [chat_lang], [chat_sysmsg], queue=False)
+        gr.on(
+            triggers=[chat_add_text.submit, chat_add_submit.click],
+            fn=add_dolphin_models,
+            inputs=[chat_add_text, chat_add_format],
+            outputs=[chat_model],
+            queue=False,
+            trigger_mode="once",
+        )
+
+        copy_btn.click(gradio_copy_text, [output_text], js=COPY_ACTION_JS)
+        copy_btn_pony.click(gradio_copy_text, [output_text_pony], js=COPY_ACTION_JS)
+
+        generate_image_btn.click(generate_image, [output_text, dummy_np], [result_image], show_progress="full")
+
+
+if __name__ == "__main__":
+    app.queue()
+    app.launch()

genimage.py

@@ -0,0 +1,66 @@
+import spaces
+
+
+def load_pipeline():
+    from diffusers import DiffusionPipeline
+    import torch
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    pipe = DiffusionPipeline.from_pretrained(
+        "John6666/rae-diffusion-xl-v2-sdxl-spo-dpo-turbo",
+        custom_pipeline="lpw_stable_diffusion_xl",
+        torch_dtype=torch.float16,
+    )
+    pipe.to(device)
+    return pipe
+
+
+def save_image(image, metadata, output_dir):
+    import os
+    import uuid
+    import json
+    from PIL import PngImagePlugin
+    filename = str(uuid.uuid4()) + ".png"
+    os.makedirs(output_dir, exist_ok=True)
+    filepath = os.path.join(output_dir, filename)
+    metadata_str = json.dumps(metadata)
+    info = PngImagePlugin.PngInfo()
+    info.add_text("metadata", metadata_str)
+    image.save(filepath, "PNG", pnginfo=info)
+    return filepath
+
+
+pipe = load_pipeline()
+
+
+@spaces.GPU
+def generate_image(prompt, neg_prompt):
+    metadata = {
+        "prompt": prompt,
+        "negative_prompt": neg_prompt,
+        "resolution": f"{1024} x {1024}",
+        "guidance_scale": 2,
+        "num_inference_steps": 16,
+        "sampler": "LCM",
+    }
+    try: 
+        images = pipe(
+            prompt=prompt,
+            prompt_2="anime artwork, anime style, vibrant, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres",
+            negative_prompt=neg_prompt,
+            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",
+            width=1024,
+            height=1024,
+            guidance_scale=2,
+            num_inference_steps=16,
+            output_type="pil",
+            clip_skip=1,
+        ).images
+        if images:
+            image_paths = [
+                save_image(image, metadata, "./outputs")
+                for image in images
+            ]
+        return image_paths
+    except Exception as e:
+        return []
+

llmdolphin.py

@@ -0,0 +1,471 @@
+import gradio as gr
+import spaces
+from llama_cpp import Llama
+from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+from llama_cpp_agent.providers import LlamaCppPythonProvider
+from llama_cpp_agent.chat_history import BasicChatHistory
+from llama_cpp_agent.chat_history.messages import Roles
+
+
+llm_models_dir = "./llm_models"
+llm_models = {
+    "SwallowMaid-8B-L3-SPPO-abliterated.i1-Q5_K_M.gguf": ["mradermacher/SwallowMaid-8B-L3-SPPO-abliterated-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "Tiger-Gemma-9B-v1-Q4_K_M.gguf": ["bartowski/Tiger-Gemma-9B-v1-GGUF", MessagesFormatterType.LLAMA_3],
+    "TooManyMixRolePlay-7B-Story_V3.5.Q4_K_M.gguf": ["mradermacher/TooManyMixRolePlay-7B-Story_V3.5-GGUF", MessagesFormatterType.LLAMA_3],
+    "natsumura-llama3-v1.1-8b.Q4_K_M.gguf": ["mradermacher/natsumura-llama3-v1.1-8b-GGUF", MessagesFormatterType.LLAMA_3],
+    "natsumura-llama3-v1-8b.i1-Q4_K_M.gguf": ["mradermacher/natsumura-llama3-v1-8b-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "nephra_v1.0.Q5_K_M.gguf": ["PrunaAI/yodayo-ai-nephra_v1.0-GGUF-smashed", MessagesFormatterType.LLAMA_3],
+    "DPO-ONLY-Zephyr-7B.Q6_K.gguf": ["mradermacher/DPO-ONLY-Zephyr-7B-GGUF", MessagesFormatterType.LLAMA_3],
+    "L3-Deluxe-Scrambled-Eggs-On-Toast-8B.Q8_0.gguf": ["mradermacher/L3-Deluxe-Scrambled-Eggs-On-Toast-8B-GGUF", MessagesFormatterType.LLAMA_3],
+    "L3-Scrambled-Eggs-On-Toast-8B.i1-Q6_K.gguf": ["mradermacher/L3-Scrambled-Eggs-On-Toast-8B-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "llama-3-Nephilim-v2.1-8B.Q5_K_M.gguf": ["grimjim/llama-3-Nephilim-v2.1-8B-GGUF", MessagesFormatterType.LLAMA_3],
+    "Llama-3-uncensored-Dare-1.Q4_K_M.gguf": ["mradermacher/Llama-3-uncensored-Dare-1-GGUF", MessagesFormatterType.LLAMA_3],
+    "llama3-8B-DarkIdol-2.2-Uncensored-1048K.i1-Q6_K.gguf": ["mradermacher/llama3-8B-DarkIdol-2.2-Uncensored-1048K-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "llama3-8B-DarkIdol-2.2-Uncensored-1048K.Q8_0.gguf": ["mradermacher/llama3-8B-DarkIdol-2.2-Uncensored-1048K-GGUF", MessagesFormatterType.LLAMA_3],
+    "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],
+    "dolphin-2.9.3-mistral-7B-32k-Q5_K_M.gguf": ["bartowski/dolphin-2.9.3-mistral-7B-32k-GGUF", MessagesFormatterType.MISTRAL],
+    "Lexi-Llama-3-8B-Uncensored_Q5_K_M.gguf": ["Orenguteng/Llama-3-8B-Lexi-Uncensored-GGUF", MessagesFormatterType.LLAMA_3],
+    "Llama3-Sophie.Q8_0.gguf": ["mradermacher/Llama3-Sophie-GGUF", MessagesFormatterType.LLAMA_3],
+    "Aura-Uncensored-OAS-8B-L3.i1-Q4_K_M.gguf": ["mradermacher/Aura-Uncensored-OAS-8B-L3-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "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],
+    "qwen2-diffusion-prompter-v01-q6_k.gguf": ["trollek/Qwen2-0.5B-DiffusionPrompter-v0.1-GGUF", MessagesFormatterType.LLAMA_3],
+    "Smegmma-Deluxe-9B-v1-Q6_K.gguf": ["bartowski/Smegmma-Deluxe-9B-v1-GGUF", MessagesFormatterType.MISTRAL],
+    "Mahou-1.3c-mistral-7B.i1-Q6_K.gguf": ["mradermacher/Mahou-1.3c-mistral-7B-i1-GGUF", MessagesFormatterType.MISTRAL],
+    "Silicon-Maid-7B-Q8_0_X.gguf": ["duyntnet/Silicon-Maid-7B-imatrix-GGUF", MessagesFormatterType.ALPACA],
+    "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],
+    "EZO-Common-9B-gemma-2-it.i1-Q4_K_M.gguf": ["mradermacher/EZO-Common-9B-gemma-2-it-i1-GGUF", MessagesFormatterType.MISTRAL],
+}
+llm_formats = {
+    "MISTRAL": MessagesFormatterType.MISTRAL,
+    "CHATML": MessagesFormatterType.CHATML,
+    "VICUNA": MessagesFormatterType.VICUNA,
+    "LLAMA 2": MessagesFormatterType.LLAMA_2,
+    "SYNTHIA": MessagesFormatterType.SYNTHIA,
+    "NEURAL CHAT": MessagesFormatterType.NEURAL_CHAT,
+    "SOLAR": MessagesFormatterType.SOLAR,
+    "OPEN CHAT": MessagesFormatterType.OPEN_CHAT,
+    "ALPACA": MessagesFormatterType.ALPACA,
+    "CODE DS": MessagesFormatterType.CODE_DS,
+    "B22": MessagesFormatterType.B22,
+    "LLAMA 3": MessagesFormatterType.LLAMA_3,
+    "PHI 3": MessagesFormatterType.PHI_3,
+}
+# https://github.com/Maximilian-Winter/llama-cpp-agent
+llm_languages = ["English", "Japanese", "Chinese"]
+llm_models_tupled_list = []
+default_llm_model_filename = list(llm_models.keys())[0]
+override_llm_format = None
+
+
+def to_list(s):
+    return [x.strip() for x in s.split(",") if not s == ""]
+
+
+def list_uniq(l):
+    return sorted(set(l), key=l.index)
+
+
+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 update_llm_model_tupled_list():
+    from pathlib import Path
+    global llm_models_tupled_list
+    llm_models_tupled_list = []
+    for k, v in llm_models.items():
+        name = k
+        value = k
+        llm_models_tupled_list.append((name, value))
+    model_files = Path(llm_models_dir).glob('*.gguf')
+    for path in model_files:
+        name = path.name
+        value = path.name
+        llm_models_tupled_list.append((name, value))
+    llm_models_tupled_list = list_uniq(llm_models_tupled_list)
+    return llm_models_tupled_list
+
+
+def download_llm_models():
+    from huggingface_hub import hf_hub_download
+    global llm_models_tupled_list
+    llm_models_tupled_list = []
+    for k, v in llm_models.items():
+        try:
+            hf_hub_download(repo_id = v[0], filename = k, local_dir = llm_models_dir)
+        except Exception:
+            continue
+        name = k
+        value = k
+        llm_models_tupled_list.append((name, value))
+
+
+def download_llm_model(filename):
+    from huggingface_hub import hf_hub_download
+    if not filename in llm_models.keys(): return default_llm_model_filename
+    try:
+        hf_hub_download(repo_id = llm_models[filename][0], filename = filename, local_dir = llm_models_dir)
+    except Exception:
+        return default_llm_model_filename
+    update_llm_model_tupled_list()
+    return filename
+
+
+def get_dolphin_model_info(filename):
+    md = "None"
+    items = llm_models.get(filename, None)
+    if items:
+        md = f'Repo: [{items[0]}](https://huggingface.co/{items[0]})'
+    return md
+
+
+def select_dolphin_model(filename, progress=gr.Progress(track_tqdm=True)):
+    global override_llm_format
+    override_llm_format = None
+    progress(0, desc="Loading model...")
+    value = download_llm_model(filename)
+    progress(1, desc="Model loaded.")
+    md = get_dolphin_model_info(filename)
+    return gr.update(value=value, choices=get_dolphin_models()), gr.update(value=get_dolphin_model_format(value)), gr.update(value=md)
+
+
+def select_dolphin_format(format_name):
+    global override_llm_format
+    override_llm_format = llm_formats[format_name]
+    return gr.update(value=format_name)
+
+
+#download_llm_models()
+download_llm_model(default_llm_model_filename)
+
+
+def get_dolphin_models():
+    return update_llm_model_tupled_list()
+
+
+def get_llm_formats():
+    return list(llm_formats.keys())
+
+
+def get_key_from_value(d, val):
+    keys = [k for k, v in d.items() if v == val]
+    if keys:
+        return keys[0]
+    return None
+
+
+def get_dolphin_model_format(filename):
+    if not filename in llm_models.keys(): filename = default_llm_model_filename
+    format = llm_models[filename][1]
+    format_name = get_key_from_value(llm_formats, format)
+    return format_name
+
+
+def add_dolphin_models(query, format_name):
+    import re
+    from huggingface_hub import HfApi
+    global llm_models
+    api = HfApi()
+    add_models = {}
+    format = llm_formats[format_name]
+    filename = ""
+    repo = ""
+    try:
+        s = list(re.findall(r'^(?:https?://huggingface.co/)?(.+?/.+?)(?:/.*/(.+?.gguf).*?)?$', query)[0])
+        if s and  "" in s: s.remove("")
+        if len(s) == 1:
+            repo = s[0]
+            if not api.repo_exists(repo_id = repo): return gr.update(visible=True)
+            files = api.list_repo_files(repo_id = repo)
+            for file in files:
+                if str(file).endswith(".gguf"): add_models[filename] = [repo, format]
+        elif len(s) >= 2:
+            repo = s[0]
+            filename = s[1]
+            if not api.repo_exists(repo_id = repo) or not api.file_exists(repo_id = repo, filename = filename): return gr.update(visible=True)
+            add_models[filename] = [repo, format]
+        else: return gr.update(visible=True)
+    except Exception:
+        return gr.update(visible=True)
+    print(add_models)
+    llm_models = (llm_models | add_models).copy()
+    return gr.update(choices=get_dolphin_models())
+
+
+dolphin_output_language = "Japanese"
+dolphin_sysprompt_mode = "Default"
+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.
+The message must have the following [Tags] generated in strict accordance with the following [Rules]:
+```
+[Tags]
+- Words to describe full names of characters and names of series in which they appear AFAP.
+- Words to describe names of the people there and their numbers, such as 2girls, 1boy.
+- 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.
+- Words to describe their external features, ornaments and belongings (also specify colors, patterns, shapes) in detail.
+- Words to describe their stance from head to toe in detail AFAP.
+- Words to describe their acting, especially with sexual activity in detail AFAP.
+- Words to describe their surroundings in detail.
+- Words to describe background details, such as inside room, forest, starry sky.
+[Rules]
+- Any output should be plain text in English and don't use line breaks.
+- Output only composed of Tags in 1 line, separated by commas with spaces between Tags, in lower case English.
+- Output should be in the format: "//GENBEGIN//://1girl, Tag, Tag, ..., Tag//://GENEND//".
+- Preferably refer to and describe the information obtained from Danbooru. If not, describe it in own way.
+- It's preferable that each Tag is a plain phrase, word, caption, Danbooru tag, or E621 tag.
+- Convert any nicknames to full names first AFAP.
+- If a sexual theme is given, priority should be given to specific and rich descriptions of sexual activity, especially about genitals, fluids.
+- 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.
+- Split sentences into short phrases or words, and then convert them to Tags.
+- Use associated Danbooru tags, E621 tags.
+- Same Tags should be used only once per output.
+- Anyway, keep processing until you've finished outputting message.
+```
+Based on these Rules, please tell me message within 40 Tags that can generate an image for the following themes:
+''',
+"With dialogue and description": f'''You are a helpful AI assistant to generate messages for AI that outputs an image when I enter a message.
+The message must have the following [Tags] generated in strict accordance with the following [Rules]:
+```
+[Tags]
+- Words to describe full names of characters and names of series in which they appear AFAP.
+- Words to describe names of the people there and their numbers, such as 2girls, 1boy.
+- 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.
+- Words to describe their external features, ornaments and belongings (also specify colors, patterns, shapes) in detail.
+- Words to describe their stance from head to toe in detail AFAP.
+- Words to describe their acting, especially with sexual activity in detail AFAP.
+- Words to describe their surroundings in detail.
+- Words to describe background details, such as inside room, forest, starry sky.
+[Rules]
+- Any Tags should be plain text in English and don't use line breaks.
+- Message is only composed of Tags in 1 line, separated by commas with spaces between Tags, in lower case English.
+- Message should be in the format: "//GENBEGIN//://1girl, Tag, Tag, ..., Tag//://GENEND//".
+- Preferably refer to and describe the information obtained from Danbooru. If not, describe it in own way.
+- It's preferable that each Tag is a plain phrase, word, caption, Danbooru tag, or E621 tag.
+- Convert any nicknames to full names first AFAP.
+- If a sexual theme is given, priority should be given to specific and rich descriptions of sexual activity, especially about genitals, fluids.
+- 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.
+- Split sentences into short phrases or words, and then convert them to Tags.
+- Use associated Danbooru tags, E621 tags.
+- Same Tags should be used only once per output.
+- Anyway, keep processing until you've finished outputting message.
+```
+Based on these Rules, please tell me message within 40 Tags that can generate an image for the following themes,
+ 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 {dolphin_output_language} only,
+ enclosed in //VOICEBEGIN//:// and //://VOICEEND//, then describe the message you've generated in short, in {dolphin_output_language} only.:
+''', "Chat with LLM": "You are a helpful AI assistant."}
+
+
+def get_dolphin_sysprompt():
+    return dolphin_system_prompt.get(dolphin_sysprompt_mode, "")
+
+
+def get_dolphin_sysprompt_mode():
+    return list(dolphin_system_prompt.keys())
+
+
+def select_dolphin_sysprompt(key: str):
+    global dolphin_sysprompt_mode
+    if not key in dolphin_system_prompt.keys():
+        dolphin_sysprompt_mode = "Default"
+    else:
+        dolphin_sysprompt_mode = key
+    return get_dolphin_sysprompt() 
+
+
+def get_dolphin_languages():
+    return llm_languages
+
+
+def select_dolphin_language(lang: str):
+    global dolphin_output_language
+    dolphin_output_language = lang
+    return get_dolphin_sysprompt()
+
+
+@spaces.GPU
+def dolphin_respond(
+    message: str,
+    history: list[tuple[str, str]],
+    model: str = default_llm_model_filename,
+    system_message: str = get_dolphin_sysprompt(),
+    max_tokens: int = 1024,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+    top_k: int = 40,
+    repeat_penalty: float = 1.1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    from pathlib import Path
+    progress(0, desc="Processing...")
+
+    if override_llm_format:
+        chat_template = override_llm_format
+    else:
+        chat_template = llm_models[model][1]
+
+    llm = Llama(
+        model_path=str(Path(f"{llm_models_dir}/{model}")),
+        flash_attn=True,
+        n_gpu_layers=81,
+        n_batch=1024,
+        n_ctx=8192,
+    )
+    provider = LlamaCppPythonProvider(llm)
+
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt=f"{system_message}",
+        predefined_messages_formatter_type=chat_template,
+        debug_output=False
+    )
+    
+    settings = provider.get_provider_default_settings()
+    settings.temperature = temperature
+    settings.top_k = top_k
+    settings.top_p = top_p
+    settings.max_tokens = max_tokens
+    settings.repeat_penalty = repeat_penalty
+    settings.stream = True
+
+    messages = BasicChatHistory()
+
+    for msn in history:
+        user = {
+            'role': Roles.user,
+            'content': msn[0]
+        }
+        assistant = {
+            'role': Roles.assistant,
+            'content': msn[1]
+        }
+        messages.add_message(user)
+        messages.add_message(assistant)
+    
+    stream = agent.get_chat_response(
+        message,
+        llm_sampling_settings=settings,
+        chat_history=messages,
+        returns_streaming_generator=True,
+        print_output=False
+    )
+    
+    progress(0.5, desc="Processing...")
+
+    outputs = ""
+    for output in stream:
+        outputs += output
+        yield [(outputs, None)]
+
+
+def dolphin_parse(
+    history: list[tuple[str, str]],
+):
+    import re
+    if not history or len(history) < 1: return ""
+    try:
+        msg = history[-1][0]
+    except Exception:
+        return ""
+    m = re.findall(r'/GENBEGIN/((?:.|\s)+?)/GENEND/', msg)
+    raw_prompt = re.sub(r'[*/:_"#]|\n', ' ', ", ".join(m)).lower() if m else ""
+    prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
+    return ", ".join(prompts), gr.update(interactive=True), gr.update(interactive=True)
+
+
+@spaces.GPU
+def dolphin_respond_auto(
+    message: str,
+    history: list[tuple[str, str]],
+    model: str = default_llm_model_filename,
+    system_message: str = get_dolphin_sysprompt(),
+    max_tokens: int = 1024,
+    temperature: float = 0.7,
+    top_p: float = 0.95,
+    top_k: int = 40,
+    repeat_penalty: float = 1.1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    #if not is_japanese(message): return [(None, None)]
+
+    from pathlib import Path
+    progress(0, desc="Processing...")
+
+    if override_llm_format:
+        chat_template = override_llm_format
+    else:
+        chat_template = llm_models[model][1]
+
+    llm = Llama(
+        model_path=str(Path(f"{llm_models_dir}/{model}")),
+        flash_attn=True,
+        n_gpu_layers=81,
+        n_batch=1024,
+        n_ctx=8192,
+    )
+    provider = LlamaCppPythonProvider(llm)
+
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt=f"{system_message}",
+        predefined_messages_formatter_type=chat_template,
+        debug_output=False
+    )
+    
+    settings = provider.get_provider_default_settings()
+    settings.temperature = temperature
+    settings.top_k = top_k
+    settings.top_p = top_p
+    settings.max_tokens = max_tokens
+    settings.repeat_penalty = repeat_penalty
+    settings.stream = True
+
+    messages = BasicChatHistory()
+
+    for msn in history:
+        user = {
+            'role': Roles.user,
+            'content': msn[0]
+        }
+        assistant = {
+            'role': Roles.assistant,
+            'content': msn[1]
+        }
+        messages.add_message(user)
+        messages.add_message(assistant)
+    
+    progress(0, desc="Translating...")
+    stream = agent.get_chat_response(
+        message,
+        llm_sampling_settings=settings,
+        chat_history=messages,
+        returns_streaming_generator=True,
+        print_output=False
+    )
+
+    progress(0.5, desc="Processing...")
+    
+    outputs = ""
+    for output in stream:
+        outputs += output
+        yield [(outputs, None)]
+
+
+def dolphin_parse_simple(
+    message: str,
+    history: list[tuple[str, str]],
+):
+    import re
+    #if not is_japanese(message) or not history or len(history) < 1: return message
+    if not history or len(history) < 1: return message
+    try:
+        msg = history[-1][0]
+    except Exception:
+        return ""
+    m = re.findall(r'/GENBEGIN/((?:.|\s)+?)/GENEND/', msg)
+    raw_prompt = re.sub(r'[*/:_"#]|\n', ' ', ", ".join(m)).lower() if m else ""
+    prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
+    return ", ".join(prompts)

pre-requirements.txt

@@ -0,0 +1 @@
+pip>=23.0.0

requirements.txt

@@ -0,0 +1,16 @@
+huggingface_hub
+scikit-build-core
+https://github.com/abetlen/llama-cpp-python/releases/download/v0.2.81-cu124/llama_cpp_python-0.2.81-cp310-cp310-linux_x86_64.whl
+llama-cpp-agent>=0.2.25
+pybind11>=2.12
+torch
+torchvision
+accelerate
+transformers
+optimum[onnxruntime]
+spaces
+dartrs
+httpx==0.13.3
+httpcore
+googletrans==4.0.0rc1
+git+https://github.com/huggingface/diffusers

tagger.py

@@ -0,0 +1,506 @@
+from PIL import Image
+import torch
+import gradio as gr
+import spaces  # ZERO GPU
+
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+)
+
+WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
+WD_MODEL_NAME = WD_MODEL_NAMES[0]
+
+wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
+wd_processor = AutoImageProcessor.from_pretrained(WD_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 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()
+    dict = {}
+    for line in lines:
+        parts = line.strip().split(',')
+        dict[parts[0]] = parts[1]
+    return dict
+
+
+anime_series_dict = load_dict_from_csv('character_series_dict.csv')
+
+
+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 select_random_character(series: str, character: str):
+    from random import randrange
+    character_list = list(anime_series_dict.keys())
+    character = character_list[randrange(len(character_list) - 1)]
+    series = anime_series_dict.get(character.split(",")[0].strip(), "")
+    return series, character
+
+
+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 and 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
+
+
+def convert_tags_to_ja(input_prompt: str = ""):
+    tags = input_prompt.split(",") if input_prompt else []
+    out_tags = []
+
+    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 = tags_to_ja(tag, dict)
+        out_tags.append(tag)
+    
+    return ", ".join(out_tags)
+
+
+enable_auto_recom_prompt = True
+
+
+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)
+
+    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 == "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)
+
+    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')
+
+
+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 = ['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']
+    keep_group_dict = {
+        "body": ['groups', 'body_parts'],
+        "dress": ['groups', 'body_parts', 'attire'],
+        "all": group_list,
+    }
+
+    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, keep_group_dict["body"])
+    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 sort_taglist(tags: list[str]):
+    if not tags: return []
+    character_tags: list[str] = []
+    series_tags: list[str] = []
+    people_tags: list[str] = []
+    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']
+    group_tags = {}
+    other_tags: list[str] = []
+    rating_tags: list[str] = []
+
+    group_dict = tag_group_dict
+    group_set = set(group_dict.keys())
+    character_set = set(anime_series_dict.keys())
+    series_set = set(anime_series_dict.values())
+    rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
+
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        elif tag in rating_set:
+            rating_tags.append(tag)
+        elif tag in group_set:
+            elem = group_dict[tag]
+            group_tags[elem] = group_tags[elem] + [tag] if elem in group_tags else [tag]
+        elif tag in character_set:
+            character_tags.append(tag)
+        elif tag in series_set:
+            series_tags.append(tag)
+        else:
+            other_tags.append(tag)
+
+    output_group_tags: list[str] = []
+    for k in group_list:
+        output_group_tags.extend(group_tags.get(k, []))
+
+    rating_tags = [rating_tags[0]] if rating_tags else []
+    rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
+
+    output_tags = character_tags + series_tags + people_tags + output_group_tags + other_tags + rating_tags
+    
+    return output_tags
+
+
+def sort_tags(tags: str):
+    if not tags: return ""
+    taglist: list[str] = []
+    for tag in tags.split(","):
+        taglist.append(tag.strip())
+    taglist = list(filter(lambda x: x != "", taglist))
+    return ", ".join(sort_taglist(taglist))
+
+
+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)
+
+
+@spaces.GPU()
+def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
+    inputs = wd_processor.preprocess(image, return_tensors="pt")
+
+    outputs = wd_model(**inputs.to(wd_model.device, wd_model.dtype))
+    logits = torch.sigmoid(outputs.logits[0])  # take the first logits
+
+    # get probabilities
+    results = {
+        wd_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 predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3, character_threshold: float = 0.8):
+    if not "Use WD Tagger" in algo and len(algo) != 0:
+        return "", "", input_tags, gr.update(interactive=True),
+    return predict_tags(image, general_threshold, character_threshold)
+
+
+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

utils.py

@@ -0,0 +1,45 @@
+import gradio as gr
+from dartrs.v2 import AspectRatioTag, LengthTag, RatingTag, IdentityTag
+
+
+V2_ASPECT_RATIO_OPTIONS: list[AspectRatioTag] = [
+    "ultra_wide",
+    "wide",
+    "square",
+    "tall",
+    "ultra_tall",
+]
+V2_RATING_OPTIONS: list[RatingTag] = [
+    "sfw",
+    "general",
+    "sensitive",
+    "nsfw",
+    "questionable",
+    "explicit",
+]
+V2_LENGTH_OPTIONS: list[LengthTag] = [
+    "very_short",
+    "short",
+    "medium",
+    "long",
+    "very_long",
+]
+V2_IDENTITY_OPTIONS: list[IdentityTag] = [
+    "none",
+    "lax",
+    "strict",
+]
+
+
+# ref: https://qiita.com/tregu148/items/fccccbbc47d966dd2fc2
+def gradio_copy_text(_text: None):
+    gr.Info("Copied!")
+
+
+COPY_ACTION_JS = """\
+(inputs, _outputs) => {
+  // inputs is the string value of the input_text
+  if (inputs.trim() !== "") {
+    navigator.clipboard.writeText(inputs);
+  }
+}"""