nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add WD14 tagger support (SwinV2 + ViT) with threshold controls

Browse Source 
Add two WD14 classifier models as alternatives to JoyCaption LLM taggers.
WD14 models are ~400MB ONNX classifiers that output Danbooru tags with
confidence scores — much faster inference, no prompt needed.

Backend:
- Add wd14-swinv2 and wd14-vit entries to MODELS dict
- Add check/download/load/generate methods for ONNX-based WD14 inference
- Add general_threshold and character_threshold params to generate_tags()
- Filter "prepend:*" tags from exclude prefixes and skip-tagged logic

Frontend:
- Add WD14 radio buttons with download status in both admin and gallery
- Toggle prompt textarea vs threshold sliders based on model selection
- Pass thresholds to API for batch and single-image tagging
- Extract _setModelStatus() helper using safe DOM methods

Dependencies: onnxruntime (optional, with clear error message)
This commit is contained in:
Vito Sansevero
2026-02-10 16:27:40 -08:00
parent 6ea8c591b7
commit 42f9c344f8
6 changed files with 511 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
web/admin.html
+49 -2
View File
@@ -980,11 +980,35 @@
<span class="text-pm-secondary text-[11px]">Checking...</span>
</div>
</label>
<label class="flex items-center justify-between p-3 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm">
<div class="flex items-center">
<input type="radio" name="autoTagModel" value="wd14-swinv2" class="w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="ml-2.5">
<span class="text-pm text-xs font-medium">WD14 SwinV2</span>
<span class="text-pm-secondary text-[11px] ml-1.5">(~400MB, fast classifier)</span>
</div>
</div>
<div id="wd14SwinV2ModelStatus" class="flex items-center gap-1.5">
<span class="text-pm-secondary text-[11px]">Checking...</span>
</div>
</label>
<label class="flex items-center justify-between p-3 bg-pm-input rounded-pm-md cursor-pointer hover:bg-pm-hover border border-pm">
<div class="flex items-center">
<input type="radio" name="autoTagModel" value="wd14-vit" class="w-3.5 h-3.5 accent-[var(--pm-accent)]">
<div class="ml-2.5">
<span class="text-pm text-xs font-medium">WD14 ViT</span>
<span class="text-pm-secondary text-[11px] ml-1.5">(~400MB, ViT architecture)</span>
</div>
</div>
<div id="wd14VitModelStatus" class="flex items-center gap-1.5">
<span class="text-pm-secondary text-[11px]">Checking...</span>
</div>
</label>
</div>
</div>
<!-- Prompt Template -->
<div class="mb-4">
<!-- Prompt Template (LLM models only) -->
<div id="promptSection" class="mb-4">
<label class="block text-xs font-medium text-pm-secondary mb-1">Tag Generation Prompt</label>
<textarea id="autoTagPrompt" rows="6"
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] font-mono resize-none focus:outline-none"
@@ -994,6 +1018,29 @@ Prioritize visual accuracy.
No sentences, no commentary, no captions. Only tags. Keep length short but descriptive.</textarea>
</div>
<!-- WD14 Threshold Controls -->
<div id="wd14ThresholdSection" class="mb-4 hidden">
<label class="block text-xs font-medium text-pm-secondary mb-2">Tag Confidence Thresholds</label>
<div class="space-y-3">
<div class="p-2.5 bg-pm-input rounded-pm-md border border-pm">
<div class="flex items-center justify-between mb-1">
<span class="text-pm text-xs">General tags</span>
<span id="wd14GeneralThresholdValue" class="text-pm-accent text-xs font-mono">0.35</span>
</div>
<input type="range" id="wd14GeneralThreshold" min="0" max="1" step="0.05" value="0.35"
class="w-full h-1.5 bg-pm-primary rounded-lg appearance-none cursor-pointer accent-[var(--pm-accent)]">
</div>
<div class="p-2.5 bg-pm-input rounded-pm-md border border-pm">
<div class="flex items-center justify-between mb-1">
<span class="text-pm text-xs">Character tags</span>
<span id="wd14CharacterThresholdValue" class="text-pm-accent text-xs font-mono">0.85</span>
</div>
<input type="range" id="wd14CharacterThreshold" min="0" max="1" step="0.05" value="0.85"
class="w-full h-1.5 bg-pm-primary rounded-lg appearance-none cursor-pointer accent-[var(--pm-accent)]">
</div>
</div>
</div>
<!-- Tagged Images Handling -->
<div class="mb-4">
<label class="block text-xs font-medium text-pm-secondary mb-2">Already Tagged Images</label>