Super-squash branch 'main' using huggingface_hub

.gitattributes

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

README.md

@@ -0,0 +1,13 @@
+---
+title: Natural Text to Danbooru Tags with Transformer V2
+emoji: 👀📦
+colorFrom: red
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.36.1
+app_file: app.py
+pinned: false
+license: openrail
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,140 @@
+import gradio as gr
+
+from v2 import (
+    V2UI,
+    parse_upsampling_output,
+    V2_ALL_MODELS,
+)
+from utils import (
+    gradio_copy_text,
+    COPY_ACTION_JS,
+    V2_ASPECT_RATIO_OPTIONS,
+    V2_RATING_OPTIONS,
+    V2_LENGTH_OPTIONS,
+    V2_IDENTITY_OPTIONS
+)
+from tagger import (
+    predict_tags,
+    convert_danbooru_to_e621_prompt,
+    remove_specific_prompt,
+    insert_recom_prompt,
+    compose_prompt_to_copy,
+    translate_prompt,
+)
+from t2t import predict_text_to_tags
+
+
+def description_ui():
+    gr.Markdown(
+        """
+## Natural Text to Danbooru Tags with Danbooru Tags Transformer V2
+Natural text => Prompt => Upsampled longer prompt
+- Mod of [ooferdoodles/text2tags-demo](https://huggingface.co/spaces/ooferdoodles/text2tags-demo) and p1atdev's [Danbooru Tags Transformer V2 Demo](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2)
+- It's buggy but seems to work for now.
+"""
+    )
+
+
+def main():
+
+    v2 = V2UI()
+
+    with gr.Blocks() as ui:
+        description_ui()
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Group():
+                    t2t_text = gr.TextArea(label="Natural text", lines=6, placeholder="Minato Aqua from hololive with pink and blue twintails in a blue maid outfit ...", value="", show_copy_button=True)
+                    with gr.Accordion(label="Advanced options", open=False):
+                        translate_t2t_text_button = gr.Button(value="Translate text to English", size="sm", variant="secondary")
+                        t2t_max_tokens = gr.Slider(0, 256, step=16, value=128, label='max_tokens')
+                        t2t_temperature = gr.Slider(0.001, 2, step=0.1, value=0.7, label='temperature')
+                        t2t_top_k = gr.Slider(0, 100, step=5, value=20, label='top_k')
+                        t2t_top_p = gr.Slider(0, 2, step=0.05, value=0.95, label='top_p')
+                        t2t_repeat_penalty = gr.Slider(0, 5, step=0.1, value=1.1, label='repeat_penalty')
+                        t2t_examples = gr.Examples([
+                                ["Minato Aqua from hololive with pink and blue twintails in a blue maid outfit"],
+                            ],
+                            t2t_text,
+                            cache_examples=False,
+                        )
+                    generate_from_text_btn = gr.Button(value="GENERATE TAGS FROM TEXT", size="lg", variant="primary")
+                
+                with gr.Group():
+                    input_character = gr.Textbox(label="Character tags", placeholder="hatsune miku", visible=False)
+                    input_copyright = gr.Textbox(label="Copyright tags", placeholder="vocaloid", visible=False)
+                    input_general = gr.TextArea(label="General tags", lines=6, placeholder="1girl, ...", value="", show_copy_button=True)
+                    input_tags_to_copy = gr.Textbox(value="", visible=False)
+                    copy_input_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False, visible=False)
+                    tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False)
+                    input_rating = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="explicit")
+                    with gr.Accordion(label="Advanced options", open=False):
+                        input_aspect_ratio = gr.Radio(label="Aspect ratio", info="The aspect ratio of the image.", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square")
+                        input_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="very_long")
+                        input_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")                    
+                        input_ban_tags = gr.Textbox(label="Ban tags", info="Tags to ban from the output.", placeholder="alternate costumen, ...", value="censored")
+                        model_name = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0])
+                        dummy_np = gr.Textbox(label="Negative prompt", value="", visible=False)
+                        recom_animagine = gr.Textbox(label="Animagine reccomended prompt", value="Animagine", visible=False)
+                        recom_pony = gr.Textbox(label="Pony reccomended prompt", value="Pony", visible=False)
+
+                generate_btn = gr.Button(value="GENERATE TAGS", size="lg", variant="primary")
+
+                with gr.Group():
+                    output_text = gr.TextArea(label="Output tags", interactive=False, show_copy_button=True)
+                    copy_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                    elapsed_time_md = gr.Markdown(label="Elapsed time", value="", visible=False)
+
+                with gr.Group():
+                    output_text_pony = gr.TextArea(label="Output tags (Pony e621 style)", interactive=False, show_copy_button=True)
+                    copy_btn_pony = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+
+        v2.input_components = [
+            model_name,
+            input_copyright,
+            input_character,
+            input_general,
+            input_rating,
+            input_aspect_ratio,
+            input_length,
+            input_identity,
+            input_ban_tags,
+        ]
+
+        translate_t2t_text_button.click(translate_prompt, inputs=[t2t_text], outputs=[t2t_text])
+
+        generate_from_text_btn.click(
+            predict_text_to_tags,
+            inputs=[t2t_text, t2t_max_tokens, t2t_temperature, t2t_top_k, t2t_top_p, t2t_repeat_penalty],
+            outputs=[
+                input_general,
+            ],
+        )
+
+        copy_input_btn.click(compose_prompt_to_copy, inputs=[input_character, input_copyright, input_general], outputs=[input_tags_to_copy]).then(
+            gradio_copy_text, inputs=[input_tags_to_copy], js=COPY_ACTION_JS,
+        )
+
+        generate_btn.click(
+            parse_upsampling_output(v2.on_generate),
+            inputs=[
+                *v2.input_components,
+            ],
+            outputs=[output_text, elapsed_time_md, copy_btn, copy_btn_pony],
+        ).then(
+            convert_danbooru_to_e621_prompt, inputs=[output_text, tag_type], outputs=[output_text_pony],
+        ).then(
+            insert_recom_prompt, inputs=[output_text, dummy_np, recom_animagine], outputs=[output_text, dummy_np],
+        ).then(
+            insert_recom_prompt, inputs=[output_text_pony, dummy_np, recom_pony], outputs=[output_text_pony, dummy_np],
+        )
+        copy_btn.click(gradio_copy_text, inputs=[output_text], js=COPY_ACTION_JS)
+        copy_btn_pony.click(gradio_copy_text, inputs=[output_text_pony], js=COPY_ACTION_JS)
+
+    ui.launch()
+
+
+if __name__ == "__main__":
+    main()
+

myt2tmod.py

@@ -0,0 +1,116 @@
+import tempfile
+import os
+import re
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import editdistance
+
+
+class TaggerLlama:
+    MODEL_URL = "John6666/llama-tagger-HF-GPTQ-4bits"
+    SAVE_NAME = "llama-tagger-HF-GPTQ-4bits"
+    TAGS_FILE_NAME = "tags.txt"
+    model = None
+    tokenizer = None
+
+    def __init__(
+        self,
+    ):
+        self.download_model()
+        self.tag_list = self.load_tags()
+
+    def download_model(
+        self,
+        model_url=None,
+        save_name=None,
+    ):
+        model_url = model_url or self.MODEL_URL  # Use self.MODEL_URL
+        save_name = save_name or self.SAVE_NAME
+
+        save_path = os.path.join(tempfile.gettempdir(), save_name)
+        if os.path.exists(save_path):
+            print("Model already exists. Skipping download.")
+            return
+        print("Downloading Model")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_URL)
+        self.model = AutoModelForCausalLM.from_pretrained(self.MODEL_URL, device_map="cuda:0")
+        self.tokenizer.save_pretrained(self.SAVE_NAME)
+        self.model.save_pretrained(self.SAVE_NAME)
+        print("Model Downloaded")
+
+    def load_tags(self):
+        module_path = os.path.abspath(__file__)
+        lookups_dir = os.path.join(os.path.dirname(module_path), "tags.txt")
+        try:
+            tags_file = lookups_dir
+            with open(tags_file, "r") as f:
+                tag_dict = [line.strip() for line in f]
+            return tag_dict
+        except IOError as e:
+            print(f"Error loading tag dictionary: {e}")
+            return []
+
+    def preprocess_tag(self, tag):
+        tag = tag.lower()
+        match = re.match(r"^([^()]*\([^()]*\))\s*.*$", tag)
+        return match.group(1) if match else tag
+
+    def find_closest_tag(self, tag, threshold, tag_list, cache={}):
+        if tag in cache:
+            return cache[tag]
+
+        closest_tag = min(tag_list, key=lambda x: editdistance.eval(tag, x))
+        if editdistance.eval(tag, closest_tag) <= threshold:
+            cache[tag] = closest_tag
+            return closest_tag
+        else:
+            return None
+
+    def correct_tags(self, tags, tag_list, preprocess=True):
+        if preprocess:
+            tags = (self.preprocess_tag(x) for x in tags)
+        corrected_tags = set()
+        for tag in tags:
+            threshold = max(1, len(tag) - 10)
+            closest_tag = self.find_closest_tag(tag, threshold, tag_list)
+            if closest_tag:
+                corrected_tags.add(closest_tag)
+        return sorted(list(corrected_tags))
+
+    def predict_tags(
+        self,
+        prompt: str,
+        max_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        repeat_penalty: float = 1.1,
+        top_k: int = 40,
+    ):
+        prompt = f"### Caption:{prompt}\n### Tags:"
+
+        input_ids = self.tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+        terminators = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("\n"),
+                       self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("### Tags:")]
+
+        raw_output = self.model.generate(
+            input_ids.to(self.model.device),
+            tokenizer=self.tokenizer,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repeat_penalty,
+            top_k=top_k,
+            do_sample=True,
+            stop_strings=["\n", "### Tags:"],
+            eos_token_id=terminators,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+        output = self.tokenizer.batch_decode(raw_output, skip_special_tokens=True)
+
+        raw_preds = re.sub('^.+\n### Tags:(.+?$)', '\\1', output[0])
+        pred_tags = [x.strip() for x in raw_preds.split(",")]
+        corrected_tags = self.correct_tags(pred_tags, self.tag_list)
+        return corrected_tags
+    
+# https://github.com/ooferdoodles1337/text2tags-lib
+# https://huggingface.co/docs/transformers/main_classes/text_generation

originalt2t.py

@@ -0,0 +1,127 @@
+from typing import Optional, List
+import tempfile
+import os
+import re
+
+import wget
+import editdistance
+from llama_cpp import Llama
+
+
+class TaggerLlama(Llama):
+    MODEL_URL = "https://huggingface.co/ooferdoodles/llama-tagger-7b/resolve/main/llama-tagger.gguf?download=true"
+    SAVE_NAME = "llama-tagger.gguf"
+    TAGS_FILE_NAME = "tags.txt"
+
+
+    def __init__(
+        self,
+        model_path: str = None,
+        **kwargs,
+    ):
+        if model_path is None:
+            model_path = os.path.join(tempfile.gettempdir(), self.SAVE_NAME)
+            self.download_model()
+        super().__init__(model_path, **kwargs)
+        self.tag_list = self.load_tags()
+
+    def download_model(
+        self,
+        model_url=None,
+        save_name=None,
+    ):
+        model_url = model_url or self.MODEL_URL  # Use self.MODEL_URL
+        save_name = save_name or self.SAVE_NAME
+
+        save_path = os.path.join(tempfile.gettempdir(), save_name)
+        if os.path.exists(save_path):
+            print("Model already exists. Skipping download.")
+            return
+        print("Downloading Model")
+        wget.download(model_url, out=save_path)
+        print("Model Downloaded")
+
+    def load_tags(self):
+        module_path = os.path.abspath(__file__)
+        lookups_dir = os.path.join(os.path.dirname(module_path), "tags.txt")
+        try:
+            tags_file = lookups_dir
+            with open(tags_file, "r") as f:
+                tag_dict = [line.strip() for line in f]
+            return tag_dict
+        except IOError as e:
+            print(f"Error loading tag dictionary: {e}")
+            return []
+
+    def preprocess_tag(self, tag):
+        tag = tag.lower()
+        match = re.match(r"^([^()]*\([^()]*\))\s*.*$", tag)
+        return match.group(1) if match else tag
+
+    def find_closest_tag(self, tag, threshold, tag_list, cache={}):
+        if tag in cache:
+            return cache[tag]
+
+        closest_tag = min(tag_list, key=lambda x: editdistance.eval(tag, x))
+        if editdistance.eval(tag, closest_tag) <= threshold:
+            cache[tag] = closest_tag
+            return closest_tag
+        else:
+            return None
+
+    def correct_tags(self, tags, tag_list, preprocess=True):
+        if preprocess:
+            tags = (self.preprocess_tag(x) for x in tags)
+        corrected_tags = set()
+        for tag in tags:
+            threshold = max(1, len(tag) - 10)
+            closest_tag = self.find_closest_tag(tag, threshold, tag_list)
+            if closest_tag:
+                corrected_tags.add(closest_tag)
+        return sorted(list(corrected_tags))
+
+    def predict_tags(
+        self,
+        prompt: str,
+        suffix: Optional[str] = None,
+        max_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        logprobs: Optional[int] = None,
+        echo: bool = False,
+        stop: Optional[List[str]] = ["/n", "### Tags:"],
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        repeat_penalty: float = 1.1,
+        top_k: int = 40,
+        stream: bool = False,
+        tfs_z: float = 1.0,
+        mirostat_mode: int = 0,
+        mirostat_tau: float = 5.0,
+        mirostat_eta: float = 0.1,
+    ):
+        prompt = f"### Caption: {prompt}\n### Tags: "
+
+        output = self.create_completion(
+            prompt=prompt,
+            suffix=suffix,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=logprobs,
+            echo=echo,
+            stop=stop,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            repeat_penalty=repeat_penalty,
+            top_k=top_k,
+            stream=stream,
+            tfs_z=tfs_z,
+            mirostat_mode=mirostat_mode,
+            mirostat_tau=mirostat_tau,
+            mirostat_eta=mirostat_eta,
+        )
+        raw_preds = output["choices"][0]["text"]
+        pred_tags = [x.strip() for x in raw_preds.split(",")]
+        corrected_tags = self.correct_tags(pred_tags, self.tag_list)
+        return corrected_tags

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

requirements.txt

@@ -0,0 +1,11 @@
+editdistance
+transformers
+accelerate
+sentencepiece
+auto-gptq
+optimum
+httpx==0.13.3
+httpcore
+googletrans==4.0.0rc1
+optimum[onnxruntime]
+dartrs

t2t.py

@@ -0,0 +1,41 @@
+from t2tmod import TaggerLlama
+import spaces
+
+
+def translate_text(text = ""):
+    def translate_to_english(prompt):
+        import httpcore
+        setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
+        from googletrans import Translator
+        translator = Translator()
+        try:
+            translated_text = translator.translate(text, src='auto', dest='en').text
+            return translated_text
+        except Exception as e:
+            return text
+
+    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
+    
+    return translate_to_english(text) if is_japanese(text) else text
+
+
+t2t_model = TaggerLlama()
+
+@spaces.GPU()
+def predict_text_to_tags(input_text: str="", max_tokens: int=128, temperature: float=0.8, top_k: int=40, top_p: float=0.95, repeat_penalty: float=1.1):
+    text = translate_text(input_text)
+    tags = t2t_model.predict_tags(text, max_tokens=max_tokens, temperature=temperature,
+                              top_k=top_k, top_p=top_p, repeat_penalty=repeat_penalty)
+    
+    if text != input_text:
+        output = text + ', ' + ', '.join(tags).replace("_", " ")
+    else:
+        output = ', '.join(tags).replace("_", " ")
+
+    return output

t2tmod.py

@@ -0,0 +1,117 @@
+import tempfile
+import os
+import re
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import editdistance
+
+
+class TaggerLlama:
+    MODEL_URL = "John6666/llama-tagger-HF-GPTQ-4bits"
+    SAVE_NAME = "llama-tagger-HF-GPTQ-4bits"
+    TAGS_FILE_NAME = "tags.txt"
+    model = None
+    tokenizer = None
+
+    def __init__(
+        self,
+    ):
+        self.download_model()
+        self.tag_list = self.load_tags()
+
+    def download_model(
+        self,
+        model_url=None,
+        save_name=None,
+    ):
+        model_url = model_url or self.MODEL_URL  # Use self.MODEL_URL
+        save_name = save_name or self.SAVE_NAME
+
+        save_path = os.path.join(tempfile.gettempdir(), save_name)
+        if os.path.exists(save_path):
+            print("Model already exists. Skipping download.")
+            return
+        print("Downloading Model")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_URL)
+        self.model = AutoModelForCausalLM.from_pretrained(self.MODEL_URL, device_map="cuda:0")
+        self.tokenizer.save_pretrained(self.SAVE_NAME)
+        self.model.save_pretrained(self.SAVE_NAME)
+        print("Model Downloaded")
+
+    def load_tags(self):
+        module_path = os.path.abspath(__file__)
+        lookups_dir = os.path.join(os.path.dirname(module_path), "tags.txt")
+        try:
+            tags_file = lookups_dir
+            with open(tags_file, "r") as f:
+                tag_dict = [line.strip() for line in f]
+            return tag_dict
+        except IOError as e:
+            print(f"Error loading tag dictionary: {e}")
+            return []
+
+    def preprocess_tag(self, tag):
+        tag = tag.lower()
+        match = re.match(r"^([^()]*\([^()]*\))\s*.*$", tag)
+        return match.group(1) if match else tag
+
+    def find_closest_tag(self, tag, threshold, tag_list, cache={}):
+        if tag in cache:
+            return cache[tag]
+
+        closest_tag = min(tag_list, key=lambda x: editdistance.eval(tag, x))
+        if editdistance.eval(tag, closest_tag) <= threshold:
+            cache[tag] = closest_tag
+            return closest_tag
+        else:
+            return None
+
+    def correct_tags(self, tags, tag_list, preprocess=True):
+        if preprocess:
+            tags = (self.preprocess_tag(x) for x in tags)
+        corrected_tags = set()
+        for tag in tags:
+            threshold = max(1, len(tag) - 10)
+            closest_tag = self.find_closest_tag(tag, threshold, tag_list)
+            if closest_tag:
+                corrected_tags.add(closest_tag)
+        return sorted(list(corrected_tags))
+
+    def predict_tags(
+        self,
+        prompt: str,
+        max_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        repeat_penalty: float = 1.1,
+        top_k: int = 40,
+    ):
+        prompt = f"### Caption:{prompt.strip()}\n### Tags:"
+
+        input_ids = self.tokenizer.encode(prompt, add_special_tokens=True, return_tensors="pt")
+        terminators = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("\n"),
+                       self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("### Tags:")]
+
+        raw_output = self.model.generate(
+            input_ids.to(self.model.device),
+            tokenizer=self.tokenizer,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repeat_penalty,
+            top_k=top_k,
+            do_sample=True,
+            stop_strings=["\n", "### Tags:"],
+            eos_token_id=terminators,
+            pad_token_id=self.tokenizer.eos_token_id,
+        )
+        output = self.tokenizer.batch_decode(raw_output, skip_special_tokens=True)
+
+        raw_preds = re.sub('^.+\n### Tags:(.+?$)', '\\1', output[0])
+        raw_preds = ",".join(raw_preds.split(",")[2:-1]) if raw_preds.split(",")[0].strip() == "1boy" else raw_preds # to avoid a mysterious bug
+        pred_tags = [x.strip() for x in raw_preds.split(",")]
+        corrected_tags = self.correct_tags(pred_tags, self.tag_list)
+        return corrected_tags
+    
+# https://github.com/ooferdoodles1337/text2tags-lib
+# https://huggingface.co/docs/transformers/main_classes/text_generation

tagger.py

@@ -0,0 +1,450 @@
+from PIL import Image
+import torch
+import gradio as gr
+import spaces  # ZERO GPU
+
+from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
+)
+
+WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
+WD_MODEL_NAME = WD_MODEL_NAMES[0]
+
+wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
+wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+
+
+def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
+    return (
+        [f"1{noun}"]
+        + [f"{num}{noun}s" for num in range(minimum + 1, maximum + 1)]
+        + [f"{maximum+1}+{noun}s"]
+    )
+
+
+PEOPLE_TAGS = (
+    _people_tag("girl") + _people_tag("boy") + _people_tag("other") + ["no humans"]
+)
+
+
+RATING_MAP = {
+    "general": "safe",
+    "sensitive": "sensitive",
+    "questionable": "nsfw",
+    "explicit": "explicit, nsfw",
+}
+DANBOORU_TO_E621_RATING_MAP = {
+    "safe": "rating_safe",
+    "sensitive": "rating_safe",
+    "nsfw": "rating_explicit",
+    "explicit, nsfw": "rating_explicit",
+    "explicit": "rating_explicit",
+    "rating:safe": "rating_safe",
+    "rating:general": "rating_safe",
+    "rating:sensitive": "rating_safe",
+    "rating:questionable, nsfw": "rating_explicit",
+    "rating:explicit, nsfw": "rating_explicit",
+}
+
+
+def load_dict_from_csv(filename):
+    with open(filename, 'r', encoding="utf-8") as f:
+        lines = f.readlines()
+    dict = {}
+    for line in lines:
+        parts = line.strip().split(',')
+        dict[parts[0]] = parts[1]
+    return dict
+
+
+anime_series_dict = load_dict_from_csv('character_series_dict.csv')
+
+
+def character_list_to_series_list(character_list):
+    output_series_tag = []
+    series_tag = ""
+    series_dict = anime_series_dict
+    for tag in character_list:
+        series_tag = series_dict.get(tag, "")
+        if tag.endswith(")"):
+            tags = tag.split("(")
+            character_tag = "(".join(tags[:-1])
+            if character_tag.endswith(" "):
+                character_tag = character_tag[:-1]
+            series_tag = tags[-1].replace(")", "")
+
+    if series_tag:
+        output_series_tag.append(series_tag)
+
+    return output_series_tag
+
+
+def danbooru_to_e621(dtag, e621_dict):
+    def d_to_e(match, e621_dict):
+        dtag = match.group(0)
+        etag = e621_dict.get(dtag.strip().replace("_", " "), "")
+        if etag:
+            return etag
+        else:
+            return dtag
+    
+    import re
+    tag = re.sub(r'[\w ]+', lambda wrapper: d_to_e(wrapper, e621_dict), dtag, 2)
+
+    return tag
+
+
+danbooru_to_e621_dict = load_dict_from_csv('danbooru_e621.csv')
+
+
+def convert_danbooru_to_e621_prompt(input_prompt: str = "", prompt_type: str = "danbooru"):
+    if prompt_type == "danbooru": return input_prompt
+    tags = input_prompt.split(",") if input_prompt else []
+    people_tags: list[str] = []
+    other_tags: list[str] = []
+    rating_tags: list[str] = []
+
+    e621_dict = danbooru_to_e621_dict
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        tag = danbooru_to_e621(tag, e621_dict)
+        if tag in PEOPLE_TAGS:        
+            people_tags.append(tag)
+        elif tag in DANBOORU_TO_E621_RATING_MAP.keys():
+            rating_tags.append(DANBOORU_TO_E621_RATING_MAP.get(tag.replace(" ",""), ""))            
+        else:
+            other_tags.append(tag)
+
+    rating_tags = sorted(set(rating_tags), key=rating_tags.index)
+    rating_tags = [rating_tags[0]] if rating_tags else []
+    rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
+
+    output_prompt = ", ".join(people_tags + other_tags + rating_tags)
+    
+    return output_prompt
+
+
+def translate_prompt(prompt: str = ""):
+    def translate_to_english(prompt):
+        import httpcore
+        setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
+        from googletrans import Translator
+        translator = Translator()
+        try:
+            translated_prompt = translator.translate(prompt, src='auto', dest='en').text
+            return translated_prompt
+        except Exception as e:
+            return prompt
+
+    def is_japanese(s):
+        import unicodedata
+        for ch in s:
+            name = unicodedata.name(ch, "") 
+            if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
+                return True
+        return False
+
+    def to_list(s):
+        return [x.strip() for x in s.split(",")]
+    
+    prompts = to_list(prompt)
+    outputs = []
+    for p in prompts:
+        p = translate_to_english(p) if is_japanese(p) else p
+        outputs.append(p)
+
+    return ", ".join(outputs)
+
+
+def translate_prompt_to_ja(prompt: str = ""):
+    def translate_to_japanese(prompt):
+        import httpcore
+        setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')
+        from googletrans import Translator
+        translator = Translator()
+        try:
+            translated_prompt = translator.translate(prompt, src='en', dest='ja').text
+            return translated_prompt
+        except Exception as e:
+            return prompt
+
+    def is_japanese(s):
+        import unicodedata
+        for ch in s:
+            name = unicodedata.name(ch, "") 
+            if "CJK UNIFIED" in name or "HIRAGANA" in name or "KATAKANA" in name:
+                return True
+        return False
+
+    def to_list(s):
+        return [x.strip() for x in s.split(",")]
+    
+    prompts = to_list(prompt)
+    outputs = []
+    for p in prompts:
+        p = translate_to_japanese(p) if not is_japanese(p) else p
+        outputs.append(p)
+
+    return ", ".join(outputs)
+
+
+def tags_to_ja(itag, dict):
+    def t_to_j(match, dict):
+        tag = match.group(0)
+        ja = dict.get(tag.strip().replace("_", " "), "")
+        if ja:
+            return ja
+        else:
+            return tag
+    
+    import re
+    tag = re.sub(r'[\w ]+', lambda wrapper: t_to_j(wrapper, dict), itag, 2)
+
+    return tag
+
+
+def convert_tags_to_ja(input_prompt: str = ""):
+    tags = input_prompt.split(",") if input_prompt else []
+    out_tags = []
+
+    tags_to_ja_dict = load_dict_from_csv('all_tags_ja_ext.csv')
+    dict = tags_to_ja_dict
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        tag = tags_to_ja(tag, dict)
+        out_tags.append(tag)
+    
+    return ", ".join(out_tags)
+
+
+def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
+    def to_list(s):
+        return [x.strip() for x in s.split(",") if not s == ""]
+    
+    def list_sub(a, b):
+        return [e for e in a if e not in b]
+    
+    def list_uniq(l):
+        return sorted(set(l), key=l.index)
+
+    animagine_ps = to_list("anime artwork, anime style, key visual, vibrant, studio anime, highly detailed, masterpiece, best quality, very aesthetic, absurdres")
+    animagine_nps = to_list("lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]")
+    pony_ps = to_list("source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres")
+    pony_nps = to_list("source_pony, source_furry, source_cartoon, score_6, score_5, score_4, busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends")
+    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 == "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
+
+
+tag_group_dict = load_dict_from_csv('tag_group.csv')
+
+
+def remove_specific_prompt(input_prompt: str = "", keep_tags: str = "all"):
+    def is_dressed(tag):
+        import re
+        p = re.compile(r'dress|cloth|uniform|costume|vest|sweater|coat|shirt|jacket|blazer|apron|leotard|hood|sleeve|skirt|shorts|pant|loafer|ribbon|necktie|bow|collar|glove|sock|shoe|boots|wear|emblem')
+        return p.search(tag)
+
+    def is_background(tag):
+        import re
+        p = re.compile(r'background|outline|light|sky|build|day|screen|tree|city')
+        return p.search(tag)
+
+    un_tags = ['solo']
+    group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
+    keep_group_dict = {
+        "body": ['groups', 'body_parts'],
+        "dress": ['groups', 'body_parts', 'attire'],
+        "all": group_list,
+    }
+
+    def is_necessary(tag, keep_tags, group_dict):
+        if keep_tags == "all":
+            return True
+        elif tag in un_tags or group_dict.get(tag, "") in explicit_group:
+            return False
+        elif keep_tags == "body" and is_dressed(tag):
+            return False
+        elif is_background(tag):
+            return False
+        else:
+            return True
+    
+    if keep_tags == "all": return input_prompt
+    keep_group = keep_group_dict.get(keep_tags, keep_group_dict["body"])
+    explicit_group = list(set(group_list) ^ set(keep_group))
+
+    tags = input_prompt.split(",") if input_prompt else []
+    people_tags: list[str] = []
+    other_tags: list[str] = []
+
+    group_dict = tag_group_dict
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        elif is_necessary(tag, keep_tags, group_dict):
+            other_tags.append(tag)
+
+    output_prompt = ", ".join(people_tags + other_tags)
+    
+    return output_prompt
+
+
+def sort_taglist(tags: list[str]):
+    if not tags: return []
+    character_tags: list[str] = []
+    series_tags: list[str] = []
+    people_tags: list[str] = []
+    group_list = ['groups', 'body_parts', 'attire', 'posture', 'objects', 'creatures', 'locations', 'disambiguation_pages', 'commonly_misused_tags', 'phrases', 'verbs_and_gerunds', 'subjective', 'nudity', 'sex_objects', 'sex', 'sex_acts', 'image_composition', 'artistic_license', 'text', 'year_tags', 'metatags']
+    group_tags = {}
+    other_tags: list[str] = []
+    rating_tags: list[str] = []
+
+    group_dict = tag_group_dict
+    group_set = set(group_dict.keys())
+    character_set = set(anime_series_dict.keys())
+    series_set = set(anime_series_dict.values())
+    rating_set = set(DANBOORU_TO_E621_RATING_MAP.keys()) | set(DANBOORU_TO_E621_RATING_MAP.values())
+
+    for tag in tags:
+        tag = tag.strip().replace("_", " ")
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        elif tag in rating_set:
+            rating_tags.append(tag)
+        elif tag in group_set:
+            elem = group_dict[tag]
+            group_tags[elem] = group_tags[elem] + [tag] if elem in group_tags else [tag]
+        elif tag in character_set:
+            character_tags.append(tag)
+        elif tag in series_set:
+            series_tags.append(tag)
+        else:
+            other_tags.append(tag)
+
+    output_group_tags: list[str] = []
+    for k in group_list:
+        output_group_tags.extend(group_tags.get(k, []))
+
+    rating_tags = [rating_tags[0]] if rating_tags else []
+    rating_tags = ["explicit, nsfw"] if rating_tags and rating_tags[0] == "explicit" else rating_tags
+
+    output_tags = character_tags + series_tags + people_tags + output_group_tags + other_tags + rating_tags
+    
+    return output_tags
+
+
+def sort_tags(tags: str):
+    if not tags: return ""
+    taglist: list[str] = []
+    for tag in tags.split(","):
+        taglist.append(tag.strip())
+    taglist = list(filter(lambda x: x != "", taglist))
+    return ", ".join(sort_taglist(taglist))
+
+
+def postprocess_results(results: dict[str, float], general_threshold: float, character_threshold: float):
+    results = {
+        k: v for k, v in sorted(results.items(), key=lambda item: item[1], reverse=True)
+    }
+
+    rating = {}
+    character = {}
+    general = {}
+
+    for k, v in results.items():
+        if k.startswith("rating:"):
+            rating[k.replace("rating:", "")] = v
+            continue
+        elif k.startswith("character:"):
+            character[k.replace("character:", "")] = v
+            continue
+
+        general[k] = v
+
+    character = {k: v for k, v in character.items() if v >= character_threshold}
+    general = {k: v for k, v in general.items() if v >= general_threshold}
+
+    return rating, character, general
+
+
+def gen_prompt(rating: list[str], character: list[str], general: list[str]):
+    people_tags: list[str] = []
+    other_tags: list[str] = []
+    rating_tag = RATING_MAP[rating[0]]
+
+    for tag in general:
+        if tag in PEOPLE_TAGS:
+            people_tags.append(tag)
+        else:
+            other_tags.append(tag)
+
+    all_tags = people_tags + other_tags
+
+    return ", ".join(all_tags)
+
+
+@spaces.GPU()
+def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
+    inputs = wd_processor.preprocess(image, return_tensors="pt")
+
+    outputs = wd_model(**inputs.to(wd_model.device, wd_model.dtype))
+    logits = torch.sigmoid(outputs.logits[0])  # take the first logits
+
+    # get probabilities
+    results = {
+        wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
+    }
+
+    # rating, character, general
+    rating, character, general = postprocess_results(
+        results, general_threshold, character_threshold
+    )
+
+    prompt = gen_prompt(
+        list(rating.keys()), list(character.keys()), list(general.keys())
+    )
+
+    output_series_tag = ""
+    output_series_list = character_list_to_series_list(character.keys())
+    if output_series_list:
+        output_series_tag = output_series_list[0]
+    else:
+        output_series_tag = ""
+
+    return output_series_tag, ", ".join(character.keys()), prompt, gr.update(interactive=True),
+
+
+def predict_tags_wd(image: Image.Image, input_tags: str, algo: list[str], general_threshold: float = 0.3, character_threshold: float = 0.8):
+    if algo and not "Use WD Tagger" in algo:
+        return "", "", input_tags, gr.update(interactive=True),
+    return predict_tags(image, general_threshold, character_threshold)
+
+
+def compose_prompt_to_copy(character: str, series: str, general: str):
+    characters = character.split(",") if character else []
+    serieses = series.split(",") if series else []
+    generals = general.split(",") if general else []
+    tags = characters + serieses + generals
+    cprompt = ",".join(tags) if tags else ""
+    return cprompt

utils.py

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

v2.py

@@ -0,0 +1,214 @@
+import time
+import os
+import torch
+from typing import Callable
+
+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:
+    import spaces
+except ImportError:
+
+    class spaces:
+        def GPU(*args, **kwargs):
+            return lambda x: x
+
+
+from output import UpsamplingOutput
+
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+
+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, auth_token=HF_TOKEN)
+    model = model_config["class"].from_pretrained(model_name, auth_token=HF_TOKEN)
+
+    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,
+        )
+