static/js/tagger.js

/* ===== 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 = `<div class="tagger-preview-container">
            <img src="${this.inputImage.dataUrl}" class="tagger-preview-img" />
        </div>`;
    },

    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 = '<span class="tagger-no-result">无匹配角色</span>';
        } 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 = `
            <span class="tagger-label-name">${this._escapeHtml(name)}</span>
            <div class="tagger-label-bar-bg">
                <div class="tagger-label-bar-fill" style="width:${pct}%"></div>
            </div>
            <span class="tagger-label-score">${pct}%</span>
        `;
        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();
    },
};