Update application files

README.md

@@ -1,6 +1,6 @@
 ---
-title: Dart
-emoji: 👁
+title: dart
+emoji: 🏷️
 colorFrom: yellow
 colorTo: pink
 sdk: gradio

app.py

@@ -1,7 +1,62 @@
+import torch
 import gradio as gr
+from optimum.onnxruntime import ORTModelForCausalLM
+from transformers import AutoTokenizer
 
-def greet(name):
-    return "Hello " + name + "!!"
+# https://huggingface.co/collections/p1atdev/dart-v2-danbooru-tags-transformer-v2-66291115701b6fe773399b0a
+model_id = "p1atdev/dart-v2-sft"
+model = ORTModelForCausalLM.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer_with_prefix_space = AutoTokenizer.from_pretrained(model_id, add_prefix_space=True)
+
+
+# https://huggingface.co/docs/transformers/v4.44.2/en/internal/generation_utils#transformers.NoBadWordsLogitsProcessor
+def get_tokens_as_list(word_list):
+    "Converts a sequence of words into a list of tokens"
+    tokens_list = []
+    for word in word_list:
+        tokenized_word = tokenizer_with_prefix_space([word], add_special_tokens=False).input_ids[0]
+        tokens_list.append(tokenized_word)
+    return tokens_list
+
+
+def generate_tags(general_tags: str):
+    # https://huggingface.co/p1atdev/dart-v2-sft#prompt-format
+    general_tags = ",".join(tag.strip() for tag in general_tags.split(",") if tag)
+    prompt = (
+        "<|bos|>"
+        # "<copyright></copyright>"
+        # "<character></character>"
+        "<|rating:general|><|aspect_ratio:tall|><|length:long|>"
+        f"<general>{general_tags}<|identity:none|><|input_end|>"
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").input_ids
+    # bad_words_ids = get_tokens_as_list(word_list=[""])
+
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            do_sample=True,
+            temperature=1.0,
+            top_p=1.0,
+            top_k=100,
+            max_new_tokens=128,
+            num_beams=1,
+            # bad_words_ids=bad_words_ids,
+        )
+
+    return ", ".join(
+        [tag for tag in tokenizer.batch_decode(outputs[0], skip_special_tokens=True) if tag.strip() != ""]
+    )
+
+
+demo = gr.Interface(
+    fn=generate_tags,
+    inputs=gr.TextArea("1girl, black hair", lines=4),
+    outputs=gr.Textbox(show_copy_button=True),
+    clear_btn=None,
+    analytics_enabled=False,
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+# Comment out the line below for default installation on Linux
+--extra-index-url https://download.pytorch.org/whl/cpu 
+
+gradio==4.42.0
+torch
+transformers
+optimum[onnxruntime] # or optimum[onnxruntime-gpu]