Update app.py

app.py

@@ -50,7 +50,7 @@ def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: floa
     
     Returns:
         results_by_cat: Dictionary mapping each category to a list of (tag, probability) above its threshold.
-        prompt_tags_by_cat: Dictionary for prompt-style output (artist, character, general).
+        prompt_tags_by_cat: Dictionary for prompt-style output (character, general).
         all_artist_tags: All artist tags (with probabilities) regardless of threshold.
     """
     probs = 1 / (1 + np.exp(-refined_logits))
@@ -59,7 +59,8 @@ def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: floa
     category_thresholds = metadata.get("category_thresholds", {})
 
     results_by_cat = {}
-    prompt_tags_by_cat = {"artist": [], "character": [], "general": []}
+    # For prompt style, only include character and general tags (artists handled separately)
+    prompt_tags_by_cat = {"character": [], "general": []}
     all_artist_tags = []
 
     for idx, prob in enumerate(probs):
@@ -77,22 +78,29 @@ def get_tags(refined_logits: np.ndarray, metadata: dict, default_threshold: floa
 def format_prompt_tags(prompt_tags_by_cat: dict, all_artist_tags: list) -> str:
     """
     Format the tags for prompt-style output.
+    Only the top artist tag is shown (regardless of threshold), and all character and general tags are shown.
     
     Returns a comma-separated string of escaped tags.
     """
-    # Sort tags within each category by probability (descending)
+    # Always select the best artist tag from all_artist_tags, regardless of threshold.
+    best_artist_tag = None
+    if all_artist_tags:
+        best_artist = max(all_artist_tags, key=lambda item: item[1])
+        best_artist_tag = escape_tag(best_artist[0])
+    
+    # Sort character and general tags by probability (descending)
     for cat in prompt_tags_by_cat:
         prompt_tags_by_cat[cat].sort(key=lambda x: x[1], reverse=True)
     
-    artist_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("artist", [])]
     character_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("character", [])]
     general_tags = [escape_tag(tag) for tag, _ in prompt_tags_by_cat.get("general", [])]
-    prompt_tags = artist_tags + character_tags + general_tags
-
-    # Ensure at least one artist tag appears if available, even if below threshold
-    if not artist_tags and all_artist_tags:
-        best_artist_tag, _ = max(all_artist_tags, key=lambda item: item[1])
-        prompt_tags.insert(0, escape_tag(best_artist_tag))
+    
+    prompt_tags = []
+    if best_artist_tag:
+        prompt_tags.append(best_artist_tag)
+    prompt_tags.extend(character_tags)
+    prompt_tags.extend(general_tags)
+    
     return ", ".join(prompt_tags) if prompt_tags else "No tags predicted."
 
 def format_detailed_output(results_by_cat: dict, all_artist_tags: list) -> str:
@@ -145,8 +153,7 @@ with demo:
         "Upload an image, adjust the threshold, and click **Tag Image** to see predictions."
     )
     gr.Markdown(
-        "*(Note: The model predicts a large number of tags across categories like character, general, artist, etc. "
-        "You can choose a concise prompt-style output or a detailed category-wise breakdown.)*"
+        "*(Note: In prompt-style output, only the top artist tag is displayed along with all character and general tags.)*"
     )
     with gr.Row():
         with gr.Column():
@@ -162,7 +169,7 @@ with demo:
                 maximum=1.0,
                 step=0.05,
                 value=DEFAULT_THRESHOLD,
-                label="Threshold"
+                label="Default Threshold"
             )
             tag_button = gr.Button("🔍 Tag Image")
         with gr.Column():