Upload 13 files

app.py

@@ -1,33 +1,105 @@
 import gradio as gr
+from model import models
 from multit2i import (
     load_models,
-    find_model_list,
     infer_multi,
+    infer_multi_random,
     save_gallery_images,
     change_model,
     get_model_info_md,
     loaded_models,
+    get_positive_prefix,
+    get_positive_suffix,
+    get_negative_prefix,
+    get_negative_suffix,
+    get_recom_prompt_type,
+    set_recom_prompt_preset,
+    get_tag_type,
+)
+from tagger.tagger import (
+    predict_tags_wd,
+    remove_specific_prompt,
+    convert_danbooru_to_e621_prompt,
+    insert_recom_prompt,
+)
+from tagger.fl2sd3longcap import predict_tags_fl2_sd3
+from tagger.v2 import (
+    V2_ALL_MODELS,
+    v2_random_prompt,
+)
+from tagger.utils import (
+    V2_ASPECT_RATIO_OPTIONS,
+    V2_RATING_OPTIONS,
+    V2_LENGTH_OPTIONS,
+    V2_IDENTITY_OPTIONS,
 )
 
 
-models = find_model_list("Yntec", [], "", "likes", 30)
-load_models(models, 10)
+load_models(models, 5)
+#load_models(models, 20) # Fetching 20 models at the same time. default: 5
 
 
-css = """"""
+css = """
+#model_info { text-align: center; display:block; }
+"""
 
 with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     with gr.Column(): 
-        model_name = gr.Dropdown(label="Select Model", choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
-        model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]))
-        image_num = gr.Slider(label="Number of Images", minimum=1, maximum=8, value=1, step=1)
-        recom_prompt = gr.Checkbox(label="Recommended Prompt", value=True)
-        prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...")
-        run_button = gr.Button("Generate Image")
+        with gr.Accordion("Advanced settings", open=True):
+            with gr.Accordion("Recommended Prompt", open=False):
+                recom_prompt_preset = gr.Radio(label="Set Presets", choices=get_recom_prompt_type(), value="Common")
+                with gr.Row():
+                    positive_prefix = gr.CheckboxGroup(label="Use Positive Prefix", choices=get_positive_prefix(), value=[])
+                    positive_suffix = gr.CheckboxGroup(label="Use Positive Suffix", choices=get_positive_suffix(), value=["Common"])
+                    negative_prefix = gr.CheckboxGroup(label="Use Negative Prefix", choices=get_negative_prefix(), value=[], visible=False)
+                    negative_suffix = gr.CheckboxGroup(label="Use Negative Suffix", choices=get_negative_suffix(), value=["Common"], visible=False)
+            with gr.Accordion("Prompt Transformer", open=False):
+                v2_rating = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="sfw")
+                v2_aspect_ratio = gr.Radio(label="Aspect ratio", info="The aspect ratio of the image.", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square", visible=False)
+                v2_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="long")
+                v2_identity = gr.Radio(label="Keep identity", info="How strictly to keep the identity of the character or subject. If you specify the detail of subject in the prompt, you should choose `strict`. Otherwise, choose `none` or `lax`. `none` is very creative but sometimes ignores the input prompt.", choices=list(V2_IDENTITY_OPTIONS), value="lax")                    
+                v2_ban_tags = gr.Textbox(label="Ban tags", info="Tags to ban from the output.", placeholder="alternate costumen, ...", value="censored")
+                v2_tag_type = gr.Radio(label="Tag Type", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru", visible=False)
+                v2_model = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0])
+            with gr.Accordion("Model", open=True):
+                model_name = gr.Dropdown(label="Select Model", show_label=False, choices=list(loaded_models.keys()), value=list(loaded_models.keys())[0], allow_custom_value=True)
+                model_info = gr.Markdown(value=get_model_info_md(list(loaded_models.keys())[0]), elem_id="model_info")
+        with gr.Group():
+            with gr.Accordion("Prompt from Image File", open=False):
+                tagger_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
+                with gr.Accordion(label="Advanced options", open=False):
+                    tagger_general_threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
+                    tagger_character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
+                    tagger_tag_type = gr.Radio(label="Convert tags to", info="danbooru for common, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
+                    tagger_recom_prompt = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)  
+                    tagger_keep_tags = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
+                tagger_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"], label="Algorithms", value=["Use WD Tagger"])
+                tagger_generate_from_image = gr.Button(value="Generate Tags from Image")
+            with gr.Row():
+                v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2)
+                v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2)
+                random_prompt = gr.Button(value="Extend Prompt 🎲", size="sm", scale=1)
+                clear_prompt = gr.Button(value="Clear Prompt 🗑️", size="sm", scale=1)
+            prompt = gr.Text(label="Prompt", lines=1, max_lines=8, placeholder="1girl, solo, ...", show_copy_button=True)
+            neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="", visible=False)
+        with gr.Row():
+            run_button = gr.Button("Generate Image", scale=6)
+            random_button = gr.Button("Random Model 🎲", scale=3)
+            image_num = gr.Number(label="Count", minimum=1, maximum=16, value=1, step=1, interactive=True, scale=1)
         results = gr.Gallery(label="Gallery", interactive=False, show_download_button=True, show_share_button=False,
                               container=True, format="png", object_fit="contain")
         image_files = gr.Files(label="Download", interactive=False)
-        clear_results = gr.Button("Clear Gallery and Download")
+        clear_results = gr.Button("Clear Gallery / Download")
+    examples = gr.Examples(
+        examples = [
+            ["souryuu asuka langley, 1girl, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors"],
+            ["sailor moon, magical girl transformation, sparkles and ribbons, soft pastel colors, crescent moon motif, starry night sky background, shoujo manga style"],
+            ["kafuu chino, 1girl, solo"],
+            ["1girl"],
+            ["beautiful sunset"],
+        ],
+        inputs=[prompt],
+    )
     gr.Markdown(
         f"""This demo was created in reference to the following demos.
 - [Nymbo/Flood](https://huggingface.co/spaces/Nymbo/Flood).
@@ -40,13 +112,49 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer_multi,
-        inputs=[prompt, model_name, recom_prompt, image_num, results],
+        inputs=[prompt, neg_prompt, results, image_num, model_name,
+                 positive_prefix, positive_suffix, negative_prefix, negative_suffix],
+        outputs=[results],
+        queue=True,
+        show_progress="full",
+        show_api=True,
+    ).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
+    gr.on(
+        triggers=[random_button.click],
+        fn=infer_multi_random,
+        inputs=[prompt, neg_prompt, results, image_num,
+                 positive_prefix, positive_suffix, negative_prefix, negative_suffix],
         outputs=[results],
         queue=True,
         show_progress="full",
         show_api=True,
     ).success(save_gallery_images, [results], [results, image_files], queue=False, show_api=False)
+    clear_prompt.click(lambda: (None, None, None), None, [prompt, v2_series, v2_character], queue=False, show_api=False)
     clear_results.click(lambda: (None, None), None, [results, image_files], queue=False, show_api=False)
+    recom_prompt_preset.change(set_recom_prompt_preset, [recom_prompt_preset],
+     [positive_prefix, positive_suffix, negative_prefix, negative_suffix], queue=False, show_api=False)
+    random_prompt.click(
+        v2_random_prompt, [prompt, v2_series, v2_character, v2_rating, v2_aspect_ratio, v2_length,
+          v2_identity, v2_ban_tags, v2_model], [prompt, v2_series, v2_character], show_api=False,
+    ).success(
+        get_tag_type, [positive_prefix, positive_suffix, negative_prefix, negative_suffix], [v2_tag_type], queue=False, show_api=False
+    ).success(
+        convert_danbooru_to_e621_prompt, [prompt, v2_tag_type], [prompt], queue=False, show_api=False,
+    )
+    tagger_generate_from_image.click(
+        predict_tags_wd,
+        [tagger_image, prompt, tagger_algorithms, tagger_general_threshold, tagger_character_threshold],
+        [v2_series, v2_character, prompt, gr.Button(visible=False)],
+        show_api=False,
+    ).success(
+        predict_tags_fl2_sd3, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False,
+    ).success(
+        remove_specific_prompt, [prompt, tagger_keep_tags], [prompt], queue=False, show_api=False,
+    ).success(
+        convert_danbooru_to_e621_prompt, [prompt, tagger_tag_type], [prompt], queue=False, show_api=False,
+    ).success(
+        insert_recom_prompt, [prompt, neg_prompt, tagger_recom_prompt], [prompt, neg_prompt], queue=False, show_api=False,
+    )
 
 demo.queue()
 demo.launch()

model.py

@@ -0,0 +1,13 @@
+from multit2i import find_model_list
+
+
+models = find_model_list("Yntec", [], "", "likes", 100)
+
+
+# Examples:
+#models = ['yodayo-ai/kivotos-xl-2.0', 'yodayo-ai/holodayo-xl-2.1'] # specific models
+#models = find_model_list("John6666", [], "", "last_modified", 20) # John6666's latest 20 models
+#models = find_model_list("John6666", ["anime"], "", "last_modified", 20) # John6666's latest 20 models with 'anime' tag
+#models = find_model_list("John6666", [], "anime", "last_modified", 20) # John6666's latest 20 models without 'anime' tag
+#models = find_model_list("", [], "", "last_modified", 20) # latest 20 text-to-image models of huggingface
+#models = find_model_list("", [], "", "downloads", 20) # monthly most downloaded 20 text-to-image models of huggingface

multit2i.py

@@ -1,12 +1,18 @@
 import gradio as gr
 import asyncio
+from threading import RLock
 from pathlib import Path
 
 
+lock = RLock()
 loaded_models = {}
 model_info_dict = {}
 
 
+def to_list(s):
+    return [x.strip() for x in s.split(",")]
+
+
 def list_sub(a, b):
     return [e for e in a if e not in b]
 
@@ -27,7 +33,7 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
     if not sort: sort = "last_modified"
     models = []
     try:
-        model_infos = api.list_models(author=author, task="text-to-image", pipeline_tag="text-to-image",
+        model_infos = api.list_models(author=author, pipeline_tag="text-to-image",
                                        tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
     except Exception as e:
         print(f"Error: Failed to list models.")
@@ -55,7 +61,8 @@ def get_t2i_model_info_dict(repo_id: str):
     if model.private or model.gated: return info
     try:
         tags = model.tags
-    except Exception:
+    except Exception as e:
+        print(e)
         return info
     if not 'diffusers' in model.tags: return info
     if 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
@@ -63,8 +70,7 @@ def get_t2i_model_info_dict(repo_id: str):
     elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
     else: info["ver"] = "Other"
     info["url"] = f"https://huggingface.co/{repo_id}/"
-    if model.card_data and model.card_data.tags:
-        info["tags"] = model.card_data.tags
+    info["tags"] = model.card_data.tags if model.card_data and model.card_data.tags else []
     info["downloads"] = model.downloads
     info["likes"] = model.likes
     info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
@@ -101,12 +107,44 @@ def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
     return gr.update(value=output_images), gr.update(value=output_paths)
 
 
+def load_from_model(model_name: str, hf_token: str = None):
+    import httpx
+    import huggingface_hub
+    from gradio.exceptions import ModelNotFoundError
+    model_url = f"https://huggingface.co/{model_name}"
+    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
+    print(f"Fetching model from: {model_url}")
+
+    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token is not None else {}
+    response = httpx.request("GET", api_url, headers=headers)
+    if response.status_code != 200:
+        raise ModelNotFoundError(
+            f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."
+        )
+    headers["X-Wait-For-Model"] = "true"
+    client = huggingface_hub.InferenceClient(model=model_name, headers=headers, token=hf_token)
+    inputs = gr.components.Textbox(label="Input")
+    outputs = gr.components.Image(label="Output")
+    fn = client.text_to_image
+
+    def query_huggingface_inference_endpoints(*data):
+        return fn(*data)
+
+    interface_info = {
+        "fn": query_huggingface_inference_endpoints,
+        "inputs": inputs,
+        "outputs": outputs,
+        "title": model_name,
+    }
+    return gr.Interface(**interface_info)
+
+
 def load_model(model_name: str):
     global loaded_models
     global model_info_dict
     if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
-        loaded_models[model_name] = gr.load(f'models/{model_name}')
+        loaded_models[model_name] = load_from_model(model_name)
         print(f"Loaded: {model_name}")
     except Exception as e:
         if model_name in loaded_models.keys(): del loaded_models[model_name]
@@ -115,8 +153,10 @@ def load_model(model_name: str):
         return None
     try:
         model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
+        print(f"Assigned: {model_name}")
     except Exception as e:
         if model_name in model_info_dict.keys(): del model_info_dict[model_name]
+        print(f"Failed to assigned: {model_name}")
         print(e)
     return loaded_models[model_name]
 
@@ -126,21 +166,124 @@ async def async_load_models(models: list, limit: int=5):
     async def async_load_model(model: str):
         async with sem:
            try:
-               return load_model(model)
+               return await asyncio.to_thread(load_model, model)
            except Exception as e:
                print(e)
     tasks = [asyncio.create_task(async_load_model(model)) for model in models]
-    return await asyncio.wait(tasks)
+    return await asyncio.gather(*tasks, return_exceptions=True)
 
 
 def load_models(models: list, limit: int=5):
-    loop = asyncio.get_event_loop()
+    loop = asyncio.new_event_loop()
     try:
         loop.run_until_complete(async_load_models(models, limit))
     except Exception as e:
         print(e)
         pass
-    loop.close()
+    finally:
+        loop.close()
+
+
+positive_prefix = {
+    "Pony": to_list("score_9, score_8_up, score_7_up"),
+    "Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
+}
+positive_suffix = {
+    "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
+    "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
+}
+negative_prefix = {
+    "Pony": to_list("score_6, score_5, score_4"),
+    "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
+    "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
+}
+negative_suffix = {
+    "Common": to_list("lowres, (bad), bad hands, bad feet, 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 Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
+    "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
+}
+positive_all = negative_all = []
+for k, v in (positive_prefix | positive_suffix).items():
+    positive_all = positive_all + v + [s.replace("_", " ") for s in v]
+positive_all = list_uniq(positive_all)
+for k, v in (negative_prefix | negative_suffix).items():
+    negative_all = negative_all + v + [s.replace("_", " ") for s in v]
+positive_all = list_uniq(positive_all)
+
+
+def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
+    def flatten(src):
+        return [item for row in src for item in row]
+    prompts = to_list(prompt)
+    neg_prompts = to_list(neg_prompt)
+    prompts = list_sub(prompts, positive_all)
+    neg_prompts = list_sub(neg_prompts, negative_all)
+    last_empty_p = [""] if not prompts and type != "None" else []
+    last_empty_np = [""] if not neg_prompts and type != "None" else []
+    prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
+    suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
+    prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
+    suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
+    prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
+    neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
+    return prompt, neg_prompt
+
+
+recom_prompt_type = {
+    "None": ([], [], [], []),
+    "Auto": ([], [], [], []),
+    "Common": ([], ["Common"], [], ["Common"]),
+    "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
+    "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
+    "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
+    "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
+}
+
+
+enable_auto_recom_prompt = False
+def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
+    global enable_auto_recom_prompt
+    if type == "Auto":  enable_auto_recom_prompt = True
+    else: enable_auto_recom_prompt = False
+    pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
+    return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
+
+
+def set_recom_prompt_preset(type: str = "None"):
+    pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
+    return pos_pre, pos_suf, neg_pre, neg_suf
+
+
+def get_recom_prompt_type():
+    type = list(recom_prompt_type.keys())
+    type.remove("Auto")
+    return type
+
+
+def get_positive_prefix():
+    return list(positive_prefix.keys())
+
+
+def get_positive_suffix():
+    return list(positive_suffix.keys())
+
+
+def get_negative_prefix():
+    return list(negative_prefix.keys())
+
+
+def get_negative_suffix():
+    return list(negative_suffix.keys())
+
+
+def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
+    tag_type = "danbooru"
+    words = pos_pre + pos_suf + neg_pre + neg_suf
+    for word in words:
+        if "Pony" in word:
+            tag_type = "e621"
+            break
+    return tag_type
 
 
 def get_model_info_md(model_name: str):
@@ -152,29 +295,56 @@ def change_model(model_name: str):
     return get_model_info_md(model_name)
 
 
-def infer(prompt: str, model_name: str, recom_prompt: bool, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt: str, neg_prompt: str, model_name: str):
     from PIL import Image
     import random
     seed = ""
     rand = random.randint(1, 500)
     for i in range(rand):
         seed += " "
-    rprompt = ", highly detailed, masterpiece, best quality, very aesthetic, absurdres, " if recom_prompt else ""
     caption = model_name.split("/")[-1]
     try:
         model = load_model(model_name)
         if not model: return (Image.Image(), None)
-        image_path = model(prompt + rprompt + seed)
-        image = Image.open(image_path).convert('RGB')
+        image_path = model(prompt + seed)
+        image = Image.open(image_path).convert('RGBA')
     except Exception as e:
         print(e)
         return (Image.Image(), None)
     return (image, caption)
 
 
-def infer_multi(prompt: str, model_name: str, recom_prompt: bool, image_num: float, results: list, progress=gr.Progress(track_tqdm=True)):
+async def infer_multi(prompt: str, neg_prompt: str, results: list, image_num: float, model_name: str,
+                 pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)):
+    from tqdm.asyncio import tqdm_asyncio
     image_num = int(image_num)
     images = results if results else []
-    for i in range(image_num):
-        images.append(infer(prompt, model_name, recom_prompt))
+    prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
+    tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for i in range(image_num)]
+    #results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = await tqdm_asyncio.gather(*tasks)
+    if not results: results = []
+    for result in results:
+        with lock:
+            if result and result[1]: images.append(result)
         yield images
+
+
+async def infer_multi_random(prompt: str, neg_prompt: str, results: list, image_num: float, 
+                 pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], progress=gr.Progress(track_tqdm=True)):
+    from tqdm.asyncio import tqdm_asyncio
+    import random
+    image_num = int(image_num)
+    images = results if results else []
+    random.seed()
+    model_names = random.choices(list(loaded_models.keys()), k = image_num)
+    prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
+    tasks = [asyncio.to_thread(infer, prompt, neg_prompt, model_name) for model_name in model_names]
+    #results = await asyncio.gather(*tasks, return_exceptions=True)
+    results = await tqdm_asyncio.gather(*tasks)
+    if not results: results = []
+    for result in results:
+        with lock:
+            if result and result[1]: images.append(result)
+        yield images
+

requirements.txt

@@ -1 +1,12 @@
-huggingface_hub
+huggingface_hub
+torch
+torchvision
+accelerate
+transformers
+optimum[onnxruntime]
+spaces
+dartrs
+httpx==0.13.3
+httpcore
+googletrans==4.0.0rc1
+timm

tagger/fl2sd3longcap.py

@@ -0,0 +1,74 @@
+from transformers import AutoProcessor, AutoModelForCausalLM
+import spaces
+import re
+from PIL import Image 
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).eval()
+fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+
+
+def fl_modify_caption(caption: str) -> str:
+    """
+    Removes specific prefixes from captions if present, otherwise returns the original caption.
+    Args:
+        caption (str): A string containing a caption.
+    Returns:
+        str: The caption with the prefix removed if it was present, or the original caption.
+    """
+    # Define the prefixes to remove
+    prefix_substrings = [
+        ('captured from ', ''),
+        ('captured at ', '')
+    ]
+    
+    # Create a regex pattern to match any of the prefixes
+    pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
+    replacers = {opening.lower(): replacer for opening, replacer in prefix_substrings}
+    
+    # Function to replace matched prefix with its corresponding replacement
+    def replace_fn(match):
+        return replacers[match.group(0).lower()]
+    
+    # Apply the regex to the caption
+    modified_caption = re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
+    
+    # If the caption was modified, return the modified version; otherwise, return the original
+    return modified_caption if modified_caption != caption else caption
+
+
+@spaces.GPU
+def fl_run_example(image):
+    task_prompt = "<DESCRIPTION>"
+    prompt = task_prompt + "Describe this image in great detail."
+
+    # Ensure the image is in RGB mode
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    inputs = fl_processor(text=prompt, images=image, return_tensors="pt")
+    generated_ids = fl_model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3
+    )
+    generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+    return fl_modify_caption(parsed_answer["<DESCRIPTION>"])
+
+
+def predict_tags_fl2_sd3(image: Image.Image, input_tags: str, algo: list[str]):
+    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)
+    
+    if not "Use Florence-2-SD3-Long-Captioner" in algo:
+        return input_tags
+    tag_list = list_uniq(to_list(input_tags) + to_list(fl_run_example(image) + ", "))
+    tag_list.remove("")
+    return ", ".join(tag_list)

tagger/output.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class UpsamplingOutput:
+    upsampled_tags: str
+
+    copyright_tags: str
+    character_tags: str
+    general_tags: str
+    rating_tag: str
+    aspect_ratio_tag: str
+    length_tag: str
+    identity_tag: str
+
+    elapsed_time: float = 0.0

tagger/tagger.py

@@ -0,0 +1,549 @@
+from PIL import Image
+import torch
+import gradio as gr
+import spaces
+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]
+
+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 = {
+    "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
+
+
+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:
+            print(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:
+            print(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(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()
+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, 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,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);
+  }
+}"""

tagger/v2.py

@@ -0,0 +1,260 @@
+import time
+import torch
+from typing import Callable
+from pathlib import Path
+
+from dartrs.v2 import (
+    V2Model,
+    MixtralModel,
+    MistralModel,
+    compose_prompt,
+    LengthTag,
+    AspectRatioTag,
+    RatingTag,
+    IdentityTag,
+)
+from dartrs.dartrs import DartTokenizer
+from dartrs.utils import get_generation_config
+
+
+import gradio as gr
+from gradio.components import Component
+
+
+try:
+    from output import UpsamplingOutput
+except:
+    from .output import UpsamplingOutput
+
+
+V2_ALL_MODELS = {
+    "dart-v2-moe-sft": {
+        "repo": "p1atdev/dart-v2-moe-sft",
+        "type": "sft",
+        "class": MixtralModel,
+    },
+    "dart-v2-sft": {
+        "repo": "p1atdev/dart-v2-sft",
+        "type": "sft",
+        "class": MistralModel,
+    },
+}
+
+
+def prepare_models(model_config: dict):
+    model_name = model_config["repo"]
+    tokenizer = DartTokenizer.from_pretrained(model_name)
+    model = model_config["class"].from_pretrained(model_name)
+
+    return {
+        "tokenizer": tokenizer,
+        "model": model,
+    }
+
+
+def normalize_tags(tokenizer: DartTokenizer, tags: str):
+    """Just remove unk tokens."""
+    return ", ".join([tag for tag in tokenizer.tokenize(tags) if tag != "<|unk|>"])
+
+
+@torch.no_grad()
+def generate_tags(
+    model: V2Model,
+    tokenizer: DartTokenizer,
+    prompt: str,
+    ban_token_ids: list[int],
+):
+    output = model.generate(
+        get_generation_config(
+            prompt,
+            tokenizer=tokenizer,
+            temperature=1,
+            top_p=0.9,
+            top_k=100,
+            max_new_tokens=256,
+            ban_token_ids=ban_token_ids,
+        ),
+    )
+
+    return output
+
+
+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"]
+)
+
+
+def gen_prompt_text(output: UpsamplingOutput):
+    # separate people tags (e.g. 1girl)
+    people_tags = []
+    other_general_tags = []
+    
+    for tag in output.general_tags.split(","):
+        tag = tag.strip()
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        else:
+            other_general_tags.append(tag)
+
+    return ", ".join(
+        [
+            part.strip()
+            for part in [
+                *people_tags,
+                output.character_tags,
+                output.copyright_tags,
+                *other_general_tags,
+                output.upsampled_tags,
+                output.rating_tag,
+            ]
+            if part.strip() != ""
+        ]
+    )
+
+
+def elapsed_time_format(elapsed_time: float) -> str:
+    return f"Elapsed: {elapsed_time:.2f} seconds"
+
+
+def parse_upsampling_output(
+    upsampler: Callable[..., UpsamplingOutput],
+):
+    def _parse_upsampling_output(*args) -> tuple[str, str, dict]:
+        output = upsampler(*args)
+
+        return (
+            gen_prompt_text(output),
+            elapsed_time_format(output.elapsed_time),
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+        )
+
+    return _parse_upsampling_output
+
+
+class V2UI:
+    model_name: str | None = None
+    model: V2Model
+    tokenizer: DartTokenizer
+
+    input_components: list[Component] = []
+    generate_btn: gr.Button
+
+    def on_generate(
+        self,
+        model_name: str,
+        copyright_tags: str,
+        character_tags: str,
+        general_tags: str,
+        rating_tag: RatingTag,
+        aspect_ratio_tag: AspectRatioTag,
+        length_tag: LengthTag,
+        identity_tag: IdentityTag,
+        ban_tags: str,
+        *args,
+    ) -> UpsamplingOutput:
+        if self.model_name is None or self.model_name != model_name:
+            models = prepare_models(V2_ALL_MODELS[model_name])
+            self.model = models["model"]
+            self.tokenizer = models["tokenizer"]
+            self.model_name = model_name
+
+        # normalize tags
+        # copyright_tags = normalize_tags(self.tokenizer, copyright_tags)
+        # character_tags = normalize_tags(self.tokenizer, character_tags)
+        # general_tags = normalize_tags(self.tokenizer, general_tags)
+
+        ban_token_ids = self.tokenizer.encode(ban_tags.strip())
+
+        prompt = compose_prompt(
+            prompt=general_tags,
+            copyright=copyright_tags,
+            character=character_tags,
+            rating=rating_tag,
+            aspect_ratio=aspect_ratio_tag,
+            length=length_tag,
+            identity=identity_tag,
+        )
+
+        start = time.time()
+        upsampled_tags = generate_tags(
+            self.model,
+            self.tokenizer,
+            prompt,
+            ban_token_ids,
+        )
+        elapsed_time = time.time() - start
+
+        return UpsamplingOutput(
+            upsampled_tags=upsampled_tags,
+            copyright_tags=copyright_tags,
+            character_tags=character_tags,
+            general_tags=general_tags,
+            rating_tag=rating_tag,
+            aspect_ratio_tag=aspect_ratio_tag,
+            length_tag=length_tag,
+            identity_tag=identity_tag,
+            elapsed_time=elapsed_time,
+        )
+
+
+def parse_upsampling_output_simple(upsampler: UpsamplingOutput):
+    return gen_prompt_text(upsampler)
+
+
+v2 = V2UI()
+
+
+def v2_upsampling_prompt(model: str = "dart-v2-moe-sft", copyright: str = "", character: str = "",
+                          general_tags: str = "", rating: str = "nsfw", aspect_ratio: str = "square",
+                            length: str = "very_long", identity: str = "lax", ban_tags: str = "censored"):
+    raw_prompt = parse_upsampling_output_simple(v2.on_generate(model, copyright, character, general_tags,
+                                                                rating, aspect_ratio, length, identity, ban_tags))
+    return raw_prompt
+
+
+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 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 v2_random_prompt(general_tags: str = "", copyright: str = "", character: str = "", rating: str = "nsfw",
+                      aspect_ratio: str = "square", length: str = "very_long", identity: str = "lax",
+                      ban_tags: str = "censored", model: str = "dart-v2-moe-sft"):
+    if copyright == "" and character == "":
+        copyright, character = select_random_character("", "")
+    raw_prompt = v2_upsampling_prompt(model, copyright, character, general_tags, rating,
+                                       aspect_ratio, length, identity, ban_tags)
+    return raw_prompt, copyright, character