/* ===== Tagger (反推) Page ===== */
window.NAITagger = {
inputImage: null, // {dataUrl, b64}
result: null, // {tags, rating, characters, general}
MODELS: [
{ id: 'SmilingWolf/wd-swinv2-tagger-v3', name: 'SwinV2 v3' },
{ id: 'SmilingWolf/wd-convnext-tagger-v3', name: 'ConvNext v3' },
{ id: 'SmilingWolf/wd-vit-tagger-v3', name: 'ViT v3' },
{ id: 'SmilingWolf/wd-vit-large-tagger-v3', name: 'ViT Large v3' },
{ id: 'SmilingWolf/wd-eva02-large-tagger-v3', name: 'EVA02 Large v3' },
{ id: 'SmilingWolf/wd-v1-4-moat-tagger-v2', name: 'MOAT v2' },
{ id: 'SmilingWolf/wd-v1-4-swinv2-tagger-v2', name: 'SwinV2 v2' },
{ id: 'SmilingWolf/wd-v1-4-convnext-tagger-v2', name: 'ConvNext v2' },
{ id: 'SmilingWolf/wd-v1-4-convnextv2-tagger-v2', name: 'ConvNextV2 v2' },
{ id: 'SmilingWolf/wd-v1-4-vit-tagger-v2', name: 'ViT v2' },
],
async handleUpload(file) {
if (!file || !file.type.startsWith('image/')) return;
const dataUrl = await NAIUtils.fileToDataUrl(file);
this.inputImage = { dataUrl, b64: NAIUtils.dataUrlToBase64(dataUrl) };
this.renderInputPreview();
this.updateSubmitBtn();
},
clearInput() {
this.inputImage = null;
this.result = null;
this.renderInputPreview();
this.renderResult();
this.updateSubmitBtn();
},
renderInputPreview() {
const area = $('#tagger-input-preview');
const clearBtn = $('#btn-tagger-clear');
if (!this.inputImage) {
area.innerHTML = '';
area.style.display = 'none';
clearBtn.style.display = 'none';
return;
}
clearBtn.style.display = '';
area.style.display = 'block';
area.innerHTML = `
`;
},
updateSubmitBtn() {
const btn = $('#btn-tagger-submit');
btn.disabled = !this.inputImage || NAIState.generating;
},
// ── Submit ──
submit() {
if (!this.inputImage || !NAIState.wsReady) return;
if (NAIState.generating) return;
NAIState.generating = true;
const btn = $('#btn-tagger-submit');
btn.disabled = true;
btn.classList.add('generating');
btn.textContent = '推理中...';
NAITimer.start();
this.result = null;
this.renderResult();
const params = {
image: this.inputImage.b64,
model_repo: $('#tagger-model-select').value,
general_thresh: parseFloat($('#tagger-general-thresh').value),
general_mcut: $('#tagger-general-mcut').checked,
character_thresh: parseFloat($('#tagger-char-thresh').value),
character_mcut: $('#tagger-char-mcut').checked,
};
NAIWS.send({ action: 'tagger', params });
},
onResult(data) {
NAIState.generating = false;
const btn = $('#btn-tagger-submit');
btn.disabled = false;
btn.classList.remove('generating');
btn.textContent = '提交';
this.updateSubmitBtn();
NAIWS.updateStatus('Tagger推理完成 ✓', false);
NAITimer.stop(true);
this.result = {
tags: data.tags || '',
rating: data.rating || {},
characters: data.characters || {},
general: data.general || {},
};
this.renderResult();
},
onError(message) {
NAIState.generating = false;
const btn = $('#btn-tagger-submit');
btn.disabled = false;
btn.classList.remove('generating');
btn.textContent = '提交';
this.updateSubmitBtn();
NAIWS.updateStatus(`Tagger错误: ${message}`, false);
NAITimer.stop(false);
},
// ── Render result ──
renderResult() {
const tagsEl = $('#tagger-output-tags');
const ratingEl = $('#tagger-output-rating');
const charEl = $('#tagger-output-chars');
const generalEl = $('#tagger-output-general');
const sendBtn = $('#btn-tagger-send');
const resultArea = $('#tagger-result-area');
if (!this.result) {
resultArea.style.display = 'none';
sendBtn.style.display = 'none';
$('#tagger-empty-hint').style.display = 'flex';
return;
}
resultArea.style.display = '';
$('#tagger-empty-hint').style.display = 'none';
// Tags string
tagsEl.value = this.result.tags;
NAIUtils.autoResize(tagsEl);
sendBtn.style.display = this.result.tags ? '' : 'none';
// Rating
ratingEl.innerHTML = '';
const sortedRating = Object.entries(this.result.rating).sort((a, b) => b[1] - a[1]);
for (const [name, score] of sortedRating) {
ratingEl.appendChild(this._createLabelBar(name, score));
}
// Characters
charEl.innerHTML = '';
const sortedChars = Object.entries(this.result.characters).sort((a, b) => b[1] - a[1]);
if (sortedChars.length === 0) {
charEl.innerHTML = '无匹配角色';
} else {
for (const [name, score] of sortedChars) {
charEl.appendChild(this._createLabelBar(name, score));
}
}
// General tags
generalEl.innerHTML = '';
const sortedGeneral = Object.entries(this.result.general).sort((a, b) => b[1] - a[1]);
for (const [name, score] of sortedGeneral.slice(0, 50)) {
generalEl.appendChild(this._createLabelBar(name, score));
}
if (sortedGeneral.length > 50) {
const more = document.createElement('div');
more.className = 'tagger-more-hint';
more.textContent = `... 还有 ${sortedGeneral.length - 50} 个标签`;
generalEl.appendChild(more);
}
},
_createLabelBar(name, score) {
const row = document.createElement('div');
row.className = 'tagger-label-row';
const pct = Math.round(score * 100);
row.innerHTML = `
${this._escapeHtml(name)}
${pct}%
`;
return row;
},
_escapeHtml(str) {
const div = document.createElement('div');
div.textContent = str;
return div.innerHTML;
},
// ── Send to generate ──
sendToGenerate() {
if (!this.result || !this.result.tags) return;
const promptEl = $('#prompt-input');
if (promptEl) {
promptEl.value = this.result.tags;
NAIUtils.autoResize(promptEl);
}
const nav = document.querySelector('.nav-item[data-page="generate"]');
if (nav) nav.click();
},
init() {
// Upload
const input = $('#tagger-input-file');
const drop = $('#tagger-input-dropzone');
drop.addEventListener('click', () => input.click());
drop.addEventListener('dragover', e => { e.preventDefault(); drop.classList.add('dragover'); });
drop.addEventListener('dragleave', () => drop.classList.remove('dragover'));
drop.addEventListener('drop', e => {
e.preventDefault(); drop.classList.remove('dragover');
if (e.dataTransfer.files.length > 0) this.handleUpload(e.dataTransfer.files[0]);
});
input.addEventListener('change', () => {
if (input.files.length > 0) this.handleUpload(input.files[0]);
input.value = '';
});
// Buttons
$('#btn-tagger-clear').addEventListener('click', () => this.clearInput());
$('#btn-tagger-submit').addEventListener('click', () => this.submit());
$('#btn-tagger-send').addEventListener('click', () => this.sendToGenerate());
// Slider value display
NAIUtils.bindSlider('tagger-general-thresh', 'tagger-general-thresh-val', v => v.toFixed(2));
NAIUtils.bindSlider('tagger-char-thresh', 'tagger-char-thresh-val', v => v.toFixed(2));
// Paste
NAIPaste.register('#tagger-input-dropzone', 'tagger', (file) => this.handleUpload(file));
this.updateSubmitBtn();
},
};