Update app.py

app.py

@@ -16,9 +16,9 @@ Demo for the WaifuDiffusion tagger models
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
 # Dataset v3 series of models:
-VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
 SWINV2_MODEL_DSV3_REPO = "SmilingWolf/wd-swinv2-tagger-v3"
 CONV_MODEL_DSV3_REPO = "SmilingWolf/wd-convnext-tagger-v3"
+VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"
 VIT_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-large-tagger-v3"
 EVA02_LARGE_MODEL_DSV3_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
 
@@ -123,9 +123,9 @@ def main():
     predictor = Predictor()
 
     model_repos = [
-        VIT_MODEL_DSV3_REPO,
         SWINV2_MODEL_DSV3_REPO,
         CONV_MODEL_DSV3_REPO,
+        VIT_MODEL_DSV3_REPO,
         VIT_LARGE_MODEL_DSV3_REPO,
         EVA02_LARGE_MODEL_DSV3_REPO,
         # ---
@@ -177,12 +177,7 @@ def main():
                        "blank_censor",
                        "blur_censor",
                        "light_censor",
-                       "mosaic_censoring"],
-                       
-    predefined_tags2 = [
-    "big, small:medium",           # If either "big" or "small" is missing, add "medium"
-    "small hand, large hand:medium hand"  # If either "small hand" or "large hand" is missing, add "medium hand"
-    ]
+                       "mosaic_censoring"]
 
     with gr.Blocks(title=TITLE) as demo:
         gr.Markdown(f"<h1 style='text-align: center;'>{TITLE}</h1>")
@@ -213,33 +208,21 @@ def main():
                         placeholder="Add tags to filter out (e.g., winter, red, from above)",
                         lines=5
                     )
-                    custom_tags = gr.Textbox(
-                        value=", ".join(predefined_tags2),
-                        label="Custom Tags (comma-separated)",
-                        placeholder="Enter custom tags to ensure they are in the output (e.g., shy, happy:sad)",
-                        lines=3
-                    )
-                    submit = gr.Button(
-                        value="Process Images", variant="primary"
-                    )
+
+                submit = gr.Button(
+                    value="Process Images", variant="primary"
+                )
 
             with gr.Column():
                 output = gr.Textbox(label="Output", lines=10)
 
-        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags, custom_tags_input):
+        def process_images(files, model_repo, general_thresh, character_thresh, filter_tags):
             images = [Image.open(file.name) for file in files]
             results = predictor.predict(images, model_repo, general_thresh, character_thresh)
-        
+
             # Parse filter tags
             filter_set = set(tag.strip().lower() for tag in filter_tags.split(","))
-        
-            # Parse custom tags and their fallback pairs
-            fallback_tags = {}
-            for pair in custom_tags_input.split(","):
-                if ":" in pair:
-                    tag, fallback = pair.split(":")
-                    fallback_tags[tag.strip().lower()] = fallback.strip().lower()
-        
+
             # Generate formatted output
             prompts = []
             for i, (general_tags, character_tags) in enumerate(results):
@@ -250,30 +233,24 @@ def main():
                 general_part = ", ".join(
                     tag.replace('_', ' ') for tag in general_tags if tag.lower() not in filter_set
                 )
-                
-                # Check if custom tags are missing and apply fallback tags
-                all_tags = set(general_tags + character_tags)
-                for tag, fallback in fallback_tags.items():
-                    if tag not in all_tags:
-                        all_tags.add(fallback)
-        
-                # Construct the final prompt
-                final_tags = ", ".join(tag.replace('_', ' ') for tag in all_tags if tag.lower() not in filter_set)
-                prompts.append(final_tags)
         
+                # Construct the prompt based on the presence of character_part
+                if character_part:
+                    prompts.append(f"{character_part}, {general_part}")
+                else:
+                    prompts.append(general_part)
+
             # Join all prompts with blank lines
             return "\n\n".join(prompts)
 
-
         submit.click(
             process_images,
-            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags, custom_tags],
+            inputs=[image_files, model_repo, general_thresh, character_thresh, filter_tags],
             outputs=output
         )
 
-
     demo.queue(max_size=10)
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()