chore: add mistral model and bump dartrs version

app.py

@@ -14,12 +14,12 @@ from utils import (
 
 
 NORMALIZE_RATING_TAG = {
-    "<|rating:sfw|>": "",
-    "<|rating:general|>": "",
-    "<|rating:sensitive|>": "sensitive",
-    "<|rating:nsfw|>": "nsfw",
-    "<|rating:questionable|>": "nsfw",
-    "<|rating:explicit|>": "nsfw, explicit",
+    "sfw": "",
+    "general": "",
+    "sensitive": "sensitive",
+    "nsfw": "nsfw",
+    "questionable": "nsfw",
+    "explicit": "nsfw, explicit",
 }
 
 

requirements.txt

@@ -4,4 +4,5 @@ transformers==4.38.2
 optimum[onnxruntime]==1.19.1
 diffusers==0.27.2
 spaces==0.26.2
-dartrs==0.1.2
+dartrs==0.1.3
+dotenv

utils.py

@@ -25,33 +25,33 @@ IMAGE_SIZES = {
     "640x1536": (640, 1536),
 }
 
-ASPECT_RATIO_OPTIONS: dict[str, AspectRatioTag] = {
-    "ultra_wide": "<|aspect_ratio:ultra_wide|>",
-    "wide": "<|aspect_ratio:wide|>",
-    "square": "<|aspect_ratio:square|>",
-    "tall": "<|aspect_ratio:tall|>",
-    "ultra_tall": "<|aspect_ratio:ultra_tall|>",
-}
-RATING_OPTIONS: dict[str, RatingTag] = {
-    "sfw": "<|rating:sfw|>",
-    "general": "<|rating:general|>",
-    "sensitive": "<|rating:sensitive|>",
-    "nsfw": "<|rating:nsfw|>",
-    "questionable": "<|rating:questionable|>",
-    "explicit": "<|rating:explicit|>",
-}
-LENGTH_OPTIONS: dict[str, LengthTag] = {
-    "very_short": "<|length:very_short|>",
-    "short": "<|length:short|>",
-    "medium": "<|length:medium|>",
-    "long": "<|length:long|>",
-    "very_long": "<|length:very_long|>",
-}
-IDENTITY_OPTIONS: dict[str, IdentityTag] = {
-    "none": "<|identity:none|>",
-    "lax": "<|identity:lax|>",
-    "strict": "<|identity:strict|>",
-}
+ASPECT_RATIO_OPTIONS: list[AspectRatioTag] = [
+    "ultra_wide",
+    "wide",
+    "square",
+    "tall",
+    "ultra_tall",
+]
+RATING_OPTIONS: list[RatingTag] = [
+    "sfw",
+    "general",
+    "sensitive",
+    "nsfw",
+    "questionable",
+    "explicit",
+]
+LENGTH_OPTIONS: list[LengthTag] = [
+    "very_short",
+    "short",
+    "medium",
+    "long",
+    "very_long",
+]
+IDENTITY_OPTIONS: list[IdentityTag] = [
+    "none",
+    "lax",
+    "strict",
+]
 
 
 PEOPLE_TAGS = [

v2.py

@@ -1,12 +1,19 @@
 import time
 import os
-
 import torch
+import dotenv
+
+dotenv.load_dotenv()
 
 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
@@ -30,11 +37,16 @@ from utils import ASPECT_RATIO_OPTIONS, RATING_OPTIONS, LENGTH_OPTIONS, IDENTITY
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
 ALL_MODELS = {
-    "dart-v2-mixtral-160m-sft": {
-        "repo": "p1atdev/dart-v2-mixtral-160m-sft-8",
+    "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,
+    },
 }
 
 
@@ -49,21 +61,9 @@ def prepare_models(model_config: dict):
     }
 
 
-# def normalize_tags(tokenizer: PreTrainedTokenizerBase, tags: str):
-#     """Just remove unk tokens."""
-#     return ", ".join(
-#         tokenizer.batch_decode(
-#             [
-#                 token
-#                 for token in tokenizer.encode_plus(
-#                     tags.strip(),
-#                     return_tensors="pt",
-#                 ).input_ids[0]
-#                 if int(token) != tokenizer.unk_token_id
-#             ],
-#             skip_special_tokens=True,
-#         )
-#     )
+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()
@@ -102,10 +102,10 @@ class V2UI:
         copyright_tags: str,
         character_tags: str,
         general_tags: str,
-        rating_option: str,
-        aspect_ratio_option: str,
-        length_option: str,
-        identity_option: str,
+        rating_tag: RatingTag,
+        aspect_ratio_tag: AspectRatioTag,
+        length_tag: LengthTag,
+        identity_tag: IdentityTag,
         ban_tags: str,
         *args,
     ) -> UpsamplingOutput:
@@ -120,10 +120,6 @@ class V2UI:
         # character_tags = normalize_tags(self.tokenizer, character_tags)
         # general_tags = normalize_tags(self.tokenizer, general_tags)
 
-        rating_tag = RATING_OPTIONS[rating_option]
-        aspect_ratio_tag = ASPECT_RATIO_OPTIONS[aspect_ratio_option]
-        length_tag = LENGTH_OPTIONS[length_option]
-        identity_tag = IDENTITY_OPTIONS[identity_option]
         ban_token_ids = self.tokenizer.encode(ban_tags.strip())
 
         prompt = compose_prompt(
@@ -175,7 +171,7 @@ class V2UI:
 
         input_rating = gr.Radio(
             label="Rating",
-            choices=list(RATING_OPTIONS.keys()),
+            choices=list(RATING_OPTIONS),
             value="general",
         )
         input_aspect_ratio = gr.Radio(
@@ -187,13 +183,13 @@ class V2UI:
         input_length = gr.Radio(
             label="Length",
             info="The total length of the tags.",
-            choices=list(LENGTH_OPTIONS.keys()),
+            choices=list(LENGTH_OPTIONS),
             value="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(IDENTITY_OPTIONS.keys()),
+            choices=list(IDENTITY_OPTIONS),
             value="none",
         )