Upload 11 files

app.py

@@ -1,7 +1,7 @@
 import spaces
 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 tagger.utils import gradio_copy_text, COPY_ACTION_JS
+from tagger.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,
@@ -59,7 +59,8 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=css, delete_ca
                     recom_pony = gr.Textbox(label="Pony reccomended prompt", value="Pony", visible=False)
             generate_image_btn = gr.Button(value="GENERATE IMAGE", size="lg", variant="primary")
             with gr.Row():
-                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")
+                result_image = gr.Gallery(label="Generated images", columns=1, object_fit="contain", container=True, preview=True, height=512,
+                                          show_label=False, show_share_button=False, show_download_button=True, interactive=False, visible=True, format="png")
     with gr.Tab("GGUF-Playground"):
         gr.Markdown("""# Chat with lots of Models and LLMs using llama.cpp
                     This tab is copy of [CaioXapelaum/GGUF-Playground](https://huggingface.co/spaces/CaioXapelaum/GGUF-Playground).<br>

genimage.py

@@ -1,20 +1,49 @@
 import spaces
+import torch
 
 
 def load_pipeline():
-    from diffusers import AutoPipelineForText2Image
-    import torch
+    from diffusers import DiffusionPipeline
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    pipe = AutoPipelineForText2Image.from_pretrained(
+    pipe = DiffusionPipeline.from_pretrained(
         "John6666/rae-diffusion-xl-v2-sdxl-spo-pcm",
-        #custom_pipeline="lpw_stable_diffusion_xl",
-        custom_pipeline="nyanko7/sdxl_smoothed_energy_guidance",
+        custom_pipeline="lpw_stable_diffusion_xl",
+        #custom_pipeline="nyanko7/sdxl_smoothed_energy_guidance",
         torch_dtype=torch.float16,
     )
     pipe.to(device)
     return pipe
 
 
+def token_auto_concat_embeds(pipe, positive, negative):
+    max_length = pipe.tokenizer.model_max_length
+    positive_length = pipe.tokenizer(positive, return_tensors="pt").input_ids.shape[-1]
+    negative_length = pipe.tokenizer(negative, return_tensors="pt").input_ids.shape[-1]
+    
+    print(f'Token length is model maximum: {max_length}, positive length: {positive_length}, negative length: {negative_length}.')
+    if max_length < positive_length or max_length < negative_length:
+        print('Concatenated embedding.')
+        if positive_length > negative_length:
+            positive_ids = pipe.tokenizer(positive, return_tensors="pt").input_ids.to("cuda")
+            negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=positive_ids.shape[-1], return_tensors="pt").input_ids.to("cuda")
+        else:
+            negative_ids = pipe.tokenizer(negative, return_tensors="pt").input_ids.to("cuda")  
+            positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=negative_ids.shape[-1],  return_tensors="pt").input_ids.to("cuda")
+    else:
+        positive_ids = pipe.tokenizer(positive, truncation=False, padding="max_length", max_length=max_length,  return_tensors="pt").input_ids.to("cuda")
+        negative_ids = pipe.tokenizer(negative, truncation=False, padding="max_length", max_length=max_length, return_tensors="pt").input_ids.to("cuda")
+    
+    positive_concat_embeds = []
+    negative_concat_embeds = []
+    for i in range(0, positive_ids.shape[-1], max_length):
+        positive_concat_embeds.append(pipe.text_encoder(positive_ids[:, i: i + max_length])[0])
+        negative_concat_embeds.append(pipe.text_encoder(negative_ids[:, i: i + max_length])[0])
+    
+    positive_prompt_embeds = torch.cat(positive_concat_embeds, dim=1)
+    negative_prompt_embeds = torch.cat(negative_concat_embeds, dim=1)
+    return positive_prompt_embeds, negative_prompt_embeds
+
+
 def save_image(image, metadata, output_dir):
     import os
     import uuid
@@ -33,26 +62,30 @@ def save_image(image, metadata, output_dir):
 pipe = load_pipeline()
 
 
+@torch.inference_mode()
 @spaces.GPU
 def generate_image(prompt, neg_prompt):
+    prompt += ", anime, masterpiece, best quality, very aesthetic, absurdres"
+    neg_prompt += ", bad hands, bad feet, 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"
     metadata = {
-        "prompt": prompt + ", anime, masterpiece, best quality, very aesthetic, absurdres",
-        "negative_prompt": neg_prompt + ", bad hands, bad feet, 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",
+        "prompt": prompt,
+        "negative_prompt": neg_prompt,
         "resolution": f"{1024} x {1024}",
         "guidance_scale": 7.0,
         "num_inference_steps": 28,
         "sampler": "Euler",
     }
     try: 
+        #positive_embeds, negative_embeds = token_auto_concat_embeds(pipe, prompt, neg_prompt)
         images = pipe(
-            prompt=prompt + ", anime, masterpiece, best quality, very aesthetic, absurdres",
-            negative_prompt=neg_prompt + ", bad hands, bad feet, 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",
+            prompt=prompt,
+            negative_prompt=neg_prompt,
             width=1024,
             height=1024,
-            guidance_scale=7.0, seg_scale=3.0, seg_applied_layers=["mid"],
+            guidance_scale=7.0,# seg_scale=3.0, seg_applied_layers=["mid"],
             num_inference_steps=28,
             output_type="pil",
-            #clip_skip=1,
+            clip_skip=1,
         ).images
         if images:
             image_paths = [

llmdolphin.py

@@ -1,5 +1,9 @@
 import spaces
 import gradio as gr
+from pathlib import Path
+import re
+import torch
+from huggingface_hub import hf_hub_download, HfApi
 from llama_cpp import Llama
 from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
 from llama_cpp_agent.providers import LlamaCppPythonProvider
@@ -7,7 +11,6 @@ from llama_cpp_agent.chat_history import BasicChatHistory
 from llama_cpp_agent.chat_history.messages import Roles
 from ja_to_danbooru.ja_to_danbooru import jatags_to_danbooru_tags
 import wrapt_timeout_decorator
-from pathlib import Path
 from llama_cpp_agent.messages_formatter import MessagesFormatter
 from formatter import mistral_v1_formatter, mistral_v2_formatter, mistral_v3_tekken_formatter
 
@@ -19,6 +22,7 @@ llm_models = {
     #"": ["", MessagesFormatterType.OPEN_CHAT],
     #"": ["", MessagesFormatterType.CHATML],
     #"": ["", MessagesFormatterType.PHI_3],
+    #"": ["", MessagesFormatterType.GEMMA_2],
     "mn-12b-lyra-v2a1-q5_k_m.gguf": ["HalleyStarbun/MN-12B-Lyra-v2a1-Q5_K_M-GGUF", MessagesFormatterType.CHATML],
     "L3-8B-Tamamo-v1.i1-Q5_K_M.gguf": ["mradermacher/L3-8B-Tamamo-v1-i1-GGUF", MessagesFormatterType.LLAMA_3],
     "MN-Chinofun-12B-2.i1-Q4_K_M.gguf": ["mradermacher/MN-Chinofun-12B-2-i1-GGUF", MessagesFormatterType.MISTRAL],
@@ -68,6 +72,19 @@ llm_models = {
     "ChatWaifu_22B_v2.0_preview.Q4_K_S.gguf": ["mradermacher/ChatWaifu_22B_v2.0_preview-GGUF", MessagesFormatterType.MISTRAL],
     "ChatWaifu_v1.4.Q5_K_M.gguf": ["mradermacher/ChatWaifu_v1.4-GGUF", MessagesFormatterType.MISTRAL],
     "ChatWaifu_v1.3.1.Q4_K_M.gguf": ["mradermacher/ChatWaifu_v1.3.1-GGUF", MessagesFormatterType.MISTRAL],
+    "Magnum_Dark_Madness_12b.Q4_K_S.gguf": ["mradermacher/Magnum_Dark_Madness_12b-GGUF", MessagesFormatterType.MISTRAL],
+    "Magnum_Lyra_Darkness_12b.Q4_K_M.gguf": ["mradermacher/Magnum_Lyra_Darkness_12b-GGUF", MessagesFormatterType.MISTRAL],
+    "Heart_Stolen-8B-task.i1-Q4_K_M.gguf": ["mradermacher/Heart_Stolen-8B-task-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "Magnum_Backyard_Party_12b.Q4_K_M.gguf": ["mradermacher/Magnum_Backyard_Party_12b-GGUF", MessagesFormatterType.MISTRAL],
+    "Magnum_Madness-12b.Q4_K_M.gguf": ["mradermacher/Magnum_Madness-12b-GGUF", MessagesFormatterType.MISTRAL],
+    "L3.1-Moe-2x8B-v0.2.i1-Q4_K_M.gguf": ["mradermacher/L3.1-Moe-2x8B-v0.2-i1-GGUF", MessagesFormatterType.LLAMA_3],
+    "Qwen2.5-14B-Wernicke-DPO.i1-Q4_K_M.gguf": ["mradermacher/Qwen2.5-14B-Wernicke-DPO-i1-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Gemma-2-Ataraxy-v4d-9B.i1-Q4_K_M.gguf": ["mradermacher/Gemma-2-Ataraxy-v4d-9B-i1-GGUF", MessagesFormatterType.GEMMA_2],
+    "qwen2.5-14b-megamerge-pt2-q5_k_m.gguf": ["CultriX/Qwen2.5-14B-MegaMerge-pt2-Q5_K_M-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "quantqwen2-merged-16bit-q4_k_m.gguf": ["davidbzyk/QuantQwen2-merged-16bit-Q4_K_M-GGUF", MessagesFormatterType.OPEN_CHAT],
+    "Mistral-nemo-ja-rp-v0.2-Q4_K_S.gguf": ["ascktgcc/Mistral-nemo-ja-rp-v0.2-GGUF", MessagesFormatterType.MISTRAL],
+    "llama3.1-darkstorm-aspire-8b-q4_k_m.gguf": ["ZeroXClem/Llama3.1-DarkStorm-Aspire-8B-Q4_K_M-GGUF", MessagesFormatterType.LLAMA_3],
+    "llama-3-yggdrasil-astralspice-8b-q4_k_m.gguf": ["ZeroXClem/Llama-3-Yggdrasil-AstralSpice-8B-Q4_K_M-GGUF", MessagesFormatterType.LLAMA_3],
     "hermes-llama3-roleplay-1000-v2.Q5_K_M.gguf": ["mradermacher/hermes-llama3-roleplay-1000-v2-GGUF", MessagesFormatterType.LLAMA_3],
     "hermes-stheno-8B-v0.1.i1-Q5_K_M.gguf": ["mradermacher/hermes-stheno-8B-v0.1-i1-GGUF", MessagesFormatterType.LLAMA_3],
     "qwen-carpmuscle-r-v0.3.Q4_K_M.gguf": ["mradermacher/qwen-carpmuscle-r-v0.3-GGUF", MessagesFormatterType.OPEN_CHAT],
@@ -832,6 +849,7 @@ llm_languages = ["English", "Japanese", "Chinese", "Korean", "Spanish", "Portugu
 llm_models_tupled_list = []
 default_llm_model_filename = list(llm_models.keys())[0]
 override_llm_format = None
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def to_list(s):
@@ -844,7 +862,6 @@ def list_uniq(l):
 
 @wrapt_timeout_decorator.timeout(dec_timeout=3.5)
 def to_list_ja(s):
-    import re
     s = re.sub(r'[、。]', ',', s)
     return [x.strip() for x in s.split(",") if not s == ""]
 
@@ -859,7 +876,6 @@ def is_japanese(s):
 
 
 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():
@@ -876,7 +892,6 @@ def update_llm_model_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():
@@ -890,7 +905,6 @@ def download_llm_models():
 
 
 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)
@@ -951,8 +965,6 @@ def get_dolphin_model_format(filename):
 
 
 def add_dolphin_models(query, format_name):
-    import re
-    from huggingface_hub import HfApi
     global llm_models
     api = HfApi()
     add_models = {}
@@ -964,20 +976,19 @@ def add_dolphin_models(query, format_name):
         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)
+            if not api.repo_exists(repo_id = repo): return gr.update()
             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)
+            if not api.repo_exists(repo_id = repo) or not api.file_exists(repo_id = repo, filename = filename): return gr.update()
             add_models[filename] = [repo, format]
-        else: return gr.update(visible=True)
+        else: return gr.update()
     except Exception as e:
         print(e)
-        return gr.update(visible=True)
-    #print(add_models)
+        return gr.update()
     llm_models = (llm_models | add_models).copy()
     update_llm_model_tupled_list()
     choices = get_dolphin_models()
@@ -1177,7 +1188,6 @@ Output should be enclosed in //GENBEGIN//:// and //://GENEND//. The text to be g
 
 
 def get_dolphin_sysprompt():
-    import re
     prompt = re.sub('<LANGUAGE>', dolphin_output_language, dolphin_system_prompt.get(dolphin_sysprompt_mode, ""))
     return prompt
 
@@ -1207,11 +1217,11 @@ def select_dolphin_language(lang: str):
 
 @wrapt_timeout_decorator.timeout(dec_timeout=5.0)
 def get_raw_prompt(msg: str):
-    import re
     m = re.findall(r'/GENBEGIN/(.+?)/GENEND/', msg, re.DOTALL)
     return re.sub(r'[*/:_"#\n]', ' ', ", ".join(m)).lower() if m else ""
 
 
+@torch.inference_mode()
 @spaces.GPU(duration=60)
 def dolphin_respond(
     message: str,
@@ -1225,87 +1235,92 @@ def dolphin_respond(
     repeat_penalty: float = 1.1,
     progress=gr.Progress(track_tqdm=True),
 ):
-    from pathlib import Path
-    progress(0, desc="Processing...")
+    try:
+        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, # 81
+            n_batch=1024,
+            n_ctx=8192, #8192
+        )
+        provider = LlamaCppPythonProvider(llm)
+
+        agent = LlamaCppAgent(
+            provider,
+            system_prompt=f"{system_message}",
+            predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
+            custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
+            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...")
 
-    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, # 81
-        n_batch=1024,
-        n_ctx=8192, #8192
-    )
-    provider = LlamaCppPythonProvider(llm)
-
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=f"{system_message}",
-        predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
-        custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
-        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)]
+        outputs = ""
+        for output in stream:
+            outputs += output
+            yield [(outputs, None)]
+    except Exception as e:
+        print(e)
+        yield [("", None)]
 
 
 def dolphin_parse(
     history: list[tuple[str, str]],
 ):
-    if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1:
-        return "", gr.update(visible=True), gr.update(visible=True)
     try:
+        if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1:
+            return "", gr.update(), gr.update()
         msg = history[-1][0]
         raw_prompt = get_raw_prompt(msg)
-    except Exception:
-        return "", gr.update(visible=True), gr.update(visible=True)
-    prompts = []
-    if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
-        prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])
-    else:
-        prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
-    return ", ".join(prompts), gr.update(interactive=True), gr.update(interactive=True)
+        prompts = []
+        if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
+            prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])
+        else:
+            prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit"])
+        return ", ".join(prompts), gr.update(interactive=True), gr.update(interactive=True)
+    except Exception as e:
+        print(e)
+        return "", gr.update(), gr.update()
 
 
+@torch.inference_mode()
 @spaces.GPU(duration=60)
 def dolphin_respond_auto(
     message: str,
@@ -1319,94 +1334,100 @@ def dolphin_respond_auto(
     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...")
+    try:
+        #if not is_japanese(message): return [(None, None)]
+        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, # 81
+            n_batch=1024,
+            n_ctx=8192, #8192
+        )
+        provider = LlamaCppPythonProvider(llm)
+
+        agent = LlamaCppAgent(
+            provider,
+            system_prompt=f"{system_message}",
+            predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
+            custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
+            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
+        )
 
-    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, # 81
-        n_batch=1024,
-        n_ctx=8192, #8192
-    )
-    provider = LlamaCppPythonProvider(llm)
-
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=f"{system_message}",
-        predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
-        custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
-        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)]
+        progress(0.5, desc="Processing...")
+
+        outputs = ""
+        for output in stream:
+            outputs += output
+            yield [(outputs, None)], gr.update(), gr.update()
+    except Exception as e:
+        print(e)
+        yield [("", None)], gr.update(), gr.update()
 
 
 def dolphin_parse_simple(
     message: str,
     history: list[tuple[str, str]],
 ):
-    #if not is_japanese(message): return message
-    if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: return message
     try:
+        #if not is_japanese(message): return message
+        if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: return message
         msg = history[-1][0]
         raw_prompt = get_raw_prompt(msg)
-    except Exception:
+        prompts = []
+        if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
+            prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit", "rating_explicit"])
+        else:
+            prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit", "rating_explicit"])
+        return ", ".join(prompts)
+    except Exception as e:
+        print(e)
         return ""
-    prompts = []
-    if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
-        prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit", "rating_explicit"])
-    else:
-        prompts = list_uniq(to_list(raw_prompt) + ["nsfw", "explicit", "rating_explicit"])
-    return ", ".join(prompts)
 
 
 # https://huggingface.co/spaces/CaioXapelaum/GGUF-Playground
 import cv2
 cv2.setNumThreads(1)
 
+
+@torch.inference_mode()
 @spaces.GPU()
 def respond_playground(
     message,
@@ -1419,47 +1440,47 @@ def respond_playground(
     top_k,
     repeat_penalty,
 ):
-    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, # 81
-        n_batch=1024,
-        n_ctx=8192, #8192
-    )
-    provider = LlamaCppPythonProvider(llm)
-
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=f"{system_message}",
-        predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
-        custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
-        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()
-
-    # Add user and assistant messages to the history
-    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 the response
     try:
+        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, # 81
+            n_batch=1024,
+            n_ctx=8192, #8192
+        )
+        provider = LlamaCppPythonProvider(llm)
+
+        agent = LlamaCppAgent(
+            provider,
+            system_prompt=f"{system_message}",
+            predefined_messages_formatter_type=chat_template if not isinstance(chat_template, MessagesFormatter) else None,
+            custom_messages_formatter=chat_template if isinstance(chat_template, MessagesFormatter) else None,
+            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()
+
+        # Add user and assistant messages to the history
+        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 the response
         stream = agent.get_chat_response(
             message,
             llm_sampling_settings=settings,
@@ -1473,4 +1494,5 @@ def respond_playground(
             outputs += output
             yield outputs
     except Exception as e:
-        yield f"Error during response generation: {str(e)}"
+        print(e)
+        yield ""

requirements.txt

@@ -10,10 +10,9 @@ accelerate
 transformers
 optimum[onnxruntime]
 dartrs
-httpx==0.13.3
-httpcore
-googletrans==4.0.0rc1
-git+https://github.com/huggingface/diffusers
+translatepy
+diffusers
 rapidfuzz
 wrapt-timeout-decorator
-opencv-python
+opencv-python
+numpy<2

tagger/tagger.py

@@ -0,0 +1,556 @@
+import spaces
+from PIL import Image
+import torch
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from pathlib import Path
+
+
+WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
+WD_MODEL_NAME = WD_MODEL_NAMES[0]
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+default_device = device
+
+try:
+    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True).to(default_device).eval()
+    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+except Exception as e:
+    print(e)
+    wd_model = wd_processor = None
+
+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 = {
+    "sfw": "safe",
+    "general": "safe",
+    "sensitive": "sensitive",
+    "questionable": "nsfw",
+    "explicit": "explicit, nsfw",
+}
+DANBOORU_TO_E621_RATING_MAP = {
+    "sfw": "rating_safe",
+    "general": "rating_safe",
+    "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",
+}
+
+
+# https://github.com/toriato/stable-diffusion-webui-wd14-tagger/blob/a9eacb1eff904552d3012babfa28b57e1d3e295c/tagger/ui.py#L368
+kaomojis = [
+    "0_0",
+    "(o)_(o)",
+    "+_+",
+    "+_-",
+    "._.",
+    "<o>_<o>",
+    "<|>_<|>",
+    "=_=",
+    ">_<",
+    "3_3",
+    "6_9",
+    ">_o",
+    "@_@",
+    "^_^",
+    "o_o",
+    "u_u",
+    "x_x",
+    "|_|",
+    "||_||",
+]
+
+
+def replace_underline(x: str):
+    return x.strip().replace("_", " ") if x not in kaomojis else x.strip()
+
+
+def to_list(s):
+    return [x.strip() for x in s.split(",") if not s == ""]
+
+
+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):
+    dict = {}
+    if not Path(filename).exists():
+        if Path('./tagger/', filename).exists(): filename = str(Path('./tagger/', filename))
+        else: return dict
+    try:
+        with open(filename, 'r', encoding="utf-8") as f:
+            lines = f.readlines()
+    except Exception:
+        print(f"Failed to open dictionary file: {filename}")
+        return 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 seed, randrange
+    seed()
+    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(replace_underline(dtag), "")
+        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 = replace_underline(tag)
+        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
+
+
+from translatepy import Translator
+translator = Translator()
+def translate_prompt_old(prompt: str = ""):
+    def translate_to_english(input: str):
+        try:
+            output = str(translator.translate(input, 'English'))
+        except Exception as e:
+            output = input
+            print(e)
+        return output
+
+    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(input: str):
+    try:
+        output = str(translator.translate(input, 'English'))
+    except Exception as e:
+        output = input
+        print(e)
+    return output
+
+
+def translate_prompt_to_ja(prompt: str = ""):
+    def translate_to_japanese(input: str):
+        try:
+            output = str(translator.translate(input, 'Japanese'))
+        except Exception as e:
+            output = input
+            print(e)
+        return output
+
+    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(replace_underline(tag), "")
+        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 = replace_underline(tag)
+        tag = tags_to_ja(tag, dict)
+        out_tags.append(tag)
+    
+    return ", ".join(out_tags)
+
+
+enable_auto_recom_prompt = True
+
+
+animagine_ps = to_list("masterpiece, best quality, very aesthetic, absurdres")
+animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]")
+pony_ps = to_list("score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres")
+pony_nps = to_list("source_pony, 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, 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("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 = {}
+    path = 'model_dict.json' if Path('model_dict.json').exists() else './tagger/model_dict.json'
+    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 = replace_underline(tag)
+        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 = replace_underline(tag)
+        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(duration=30)
+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
+    if device != default_device: wd_model.to(device=device)
+    results = {
+        wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
+    }
+    if device != default_device: wd_model.to(device=default_device)
+    # 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, input_series: str = "", input_character: str = ""):
+    if not "Use WD Tagger" in algo and len(algo) != 0:
+        return input_series, input_character, 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

tagger/utils.py

@@ -0,0 +1,50 @@
+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);
+  }
+}"""
+
+
+def gradio_copy_prompt(prompt: str):
+    gr.Info("Copied!")
+    return prompt