Upload 6 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: WaifuDiffusion Tagger
+title: WaifuDiffusion v1.4 Tags
 emoji: 💬
 colorFrom: blue
 colorTo: red
@@ -7,6 +7,7 @@ sdk: gradio
 sdk_version: 4.20.1
 app_file: app.py
 pinned: false
+duplicated_from: SmilingWolf/wd-v1-4-tags
 ---
 
 # Configuration

app.py

@@ -8,8 +8,15 @@ import onnxruntime as rt
 import pandas as pd
 from PIL import Image
 
+
+
 TITLE = "WaifuDiffusion Tagger"
+
 DESCRIPTION = """
+This is an edited version of SmilingWolf's wd-1.4 taggs, which I have modified so that you don't have to remove the commas when you label an image for a booru website
+
+https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags
+
 Demo for the WaifuDiffusion tagger models
 
 Example image by [ほし☆☆☆](https://www.pixiv.net/en/users/43565085)
@@ -68,9 +75,9 @@ def parse_args() -> argparse.Namespace:
 
 def load_labels(dataframe) -> list[str]:
     name_series = dataframe["name"]
-    name_series = name_series.map(
-        lambda x: x.replace("_", " ") if x not in kaomojis else x
-    )
+    # name_series = name_series.map(
+    #     lambda x: x.replace("_", " ") if x not in kaomojis else x
+    # )
     tag_names = name_series.tolist()
 
     rating_indexes = list(np.where(dataframe["category"] == 9)[0])
@@ -213,11 +220,17 @@ class Predictor:
             reverse=True,
         )
         sorted_general_strings = [x[0] for x in sorted_general_strings]
+        sorted_booru_strings = (
+            " ".join(sorted_general_strings)
+        )
         sorted_general_strings = (
             ", ".join(sorted_general_strings).replace("(", "\(").replace(")", "\)")
         )
+        # sorted_general_strings = sorted_general_strings.map(
+        # lambda x: x.replace("_", " ") if x not in kaomojis else x
+        # )
 
-        return sorted_general_strings, rating, character_res, general_res
+        return sorted_general_strings, sorted_booru_strings, rating, character_res, general_res
 
 
 def main():
@@ -294,12 +307,14 @@ def main():
                         submit = gr.Button(value="Submit", variant="primary", size="lg")
                 with gr.Column(variant="panel"):
                     sorted_general_strings = gr.Textbox(label="Output (string)")
+                    sorted_booru_strings = gr.Textbox(label="Output (Booru)")
                     rating = gr.Label(label="Rating")
                     character_res = gr.Label(label="Output (characters)")
                     general_res = gr.Label(label="Output (tags)")
                     clear.add(
                         [
                             sorted_general_strings,
+                            sorted_booru_strings,
                             rating,
                             character_res,
                             general_res,
@@ -316,7 +331,7 @@ def main():
                 character_thresh,
                 character_mcut_enabled,
             ],
-            outputs=[sorted_general_strings, rating, character_res, general_res],
+            outputs=[sorted_general_strings,sorted_booru_strings, rating, character_res, general_res],
         )
 
         gr.Examples(

gitignore

@@ -0,0 +1 @@
+images

requirements.txt

@@ -1,3 +1,3 @@
 pillow>=9.0.0
 onnxruntime>=1.12.0
-huggingface-hub
+huggingface-hub