static/js/extra-input.js

/* ===== Extra Input Module (img2img / inpaint) ===== */
window.NAIExtraInput = {
    async handleUpload(file) {
        if (!file || !file.type.startsWith('image/')) return;
        const dataUrl = await NAIUtils.fileToDataUrl(file);
        NAIState.extraInputImage = { file, dataUrl, b64: NAIUtils.dataUrlToBase64(dataUrl) };
        this.renderPreview();
    },

    clear() {
        NAIState.extraInputImage = null;
        NAIState.maskHistory = [];
        NAIState.maskHistoryIndex = -1;
        NAIState.maskCanvas = null;
        NAIState.maskCtx = null;
        NAIState.overlayCanvas = null;
        NAIState.overlayCtx = null;
        this.renderPreview();
    },

    setMode(mode) {
        NAIState.extraInputMode = mode;
        if (NAIState.extraInputImage) this.renderPreview();
    },

    renderPreview() {
        const state = NAIState;
        const previewArea = $('#extra-input-preview');
        const slidersArea = $('#extra-input-sliders');
        const clearBtn = $('#btn-extra-clear');
        const U = NAIUtils;

        if (!state.extraInputImage) {
            previewArea.innerHTML = '';
            previewArea.style.display = 'none';
            slidersArea.innerHTML = '';
            clearBtn.style.display = 'none';
            return;
        }

        clearBtn.style.display = '';
        previewArea.style.display = 'block';
        previewArea.innerHTML = '';

        const mode = state.extraInputMode;
        const container = document.createElement('div');
        container.className = 'extra-preview-container';

        const img = document.createElement('img');
        img.src = state.extraInputImage.dataUrl;
        img.className = 'extra-preview-img';
        container.appendChild(img);

        if (mode === 'inpaint') {
            const overlayCanvas = document.createElement('canvas');
            overlayCanvas.className = 'extra-mask-canvas';

            const maskCanvas = document.createElement('canvas');
            maskCanvas.style.display = 'none';
            container.appendChild(overlayCanvas);
            container.appendChild(maskCanvas);

            const initCanvas = () => {
                const nw = img.naturalWidth;
                const nh = img.naturalHeight;
                overlayCanvas.width = nw;
                overlayCanvas.height = nh;
                maskCanvas.width = nw;
                maskCanvas.height = nh;

                // CSS sizing: match the displayed image exactly
                // Use the image's rendered size
                const updateCanvasSize = () => {
                    const containerRect = img.parentElement.getBoundingClientRect();
                    const imgRect = img.getBoundingClientRect();
                    // Position canvas exactly over the rendered image
                    const offsetLeft = imgRect.left - containerRect.left;
                    const offsetTop = imgRect.top - containerRect.top;
                    overlayCanvas.style.width = imgRect.width + 'px';
                    overlayCanvas.style.height = imgRect.height + 'px';
                    overlayCanvas.style.left = offsetLeft + 'px';
                    overlayCanvas.style.top = offsetTop + 'px';
                    overlayCanvas.style.transform = 'none';
                };
                updateCanvasSize();

                state.overlayCanvas = overlayCanvas;
                state.overlayCtx = overlayCanvas.getContext('2d');
                state.maskCanvas = maskCanvas;
                state.maskCtx = maskCanvas.getContext('2d');

                state.maskCtx.fillStyle = '#000000';
                state.maskCtx.fillRect(0, 0, maskCanvas.width, maskCanvas.height);

                if (state.maskHistory.length > 0 && state.maskHistoryIndex >= 0) {
                    state.maskCtx.putImageData(state.maskHistory[state.maskHistoryIndex].mask, 0, 0);
                    state.overlayCtx.putImageData(state.maskHistory[state.maskHistoryIndex].overlay, 0, 0);
                } else {
                    this.saveMaskState();
                }

                this.setupPainting(overlayCanvas);

                new ResizeObserver(() => {
                    updateCanvasSize();
                }).observe(img);
            };

            img.onload = initCanvas;
            if (img.complete && img.naturalWidth > 0) initCanvas();
        }

        previewArea.appendChild(container);

        // Sliders
        slidersArea.innerHTML = '';

        if (mode === 'img2img') {
            slidersArea.appendChild(U.createDynSlider('去噪强度', 'extra-denoise', 0, 0.99, 0.01, 0.7));
            slidersArea.appendChild(U.createDynSlider('噪声', 'extra-noise', 0, 1, 0.1, 0));
        } else {
            slidersArea.appendChild(U.createDynSlider('重绘强度', 'extra-inpaint-strength', 0.01, 1, 0.01, 1));

            // Brush size
            const brushGroup = document.createElement('div');
            brushGroup.className = 'form-group dyn-slider-group';
            brushGroup.innerHTML = `
                <label>画笔大小</label>
                <div class="slider-group">
                    <input type="range" id="extra-brush-size" min="5" max="100" step="1" value="${state.brushSize}">
                    <span id="extra-brush-size-val" class="slider-value">${state.brushSize}</span>
                </div>`;
            slidersArea.appendChild(brushGroup);
            setTimeout(() => {
                const bs = $('#extra-brush-size');
                const bsv = $('#extra-brush-size-val');
                if (bs && bsv) bs.addEventListener('input', () => { state.brushSize = parseInt(bs.value); bsv.textContent = bs.value; });
            }, 0);

            // Add original image checkbox
            const overlayGroup = document.createElement('div');
            overlayGroup.className = 'extra-checkbox-row';
            overlayGroup.innerHTML = `
                <input type="checkbox" id="extra-add-original" checked>
                <label for="extra-add-original">覆盖原图</label>
            `;
            slidersArea.appendChild(overlayGroup);

            // Undo / Clear
            const maskBtns = document.createElement('div');
            maskBtns.className = 'mask-btn-row';
            maskBtns.innerHTML = `
                <button id="btn-mask-undo" class="btn-mask">↩ 撤销</button>
                <button id="btn-mask-clear" class="btn-mask">🗑 清除遮罩</button>`;
            slidersArea.appendChild(maskBtns);
            setTimeout(() => {
                $('#btn-mask-undo').addEventListener('click', () => this.undoMask());
                $('#btn-mask-clear').addEventListener('click', () => this.clearMask());
            }, 0);
        }
    },

    // ── Mask painting ──
    setupPainting(canvas) {
        let lastX, lastY;
        const state = NAIState;

        const getCoords = (e) => {
            const rect = canvas.getBoundingClientRect();
            const scaleX = canvas.width / rect.width;
            const scaleY = canvas.height / rect.height;
            const [cx, cy] = e.touches
                ? [e.touches[0].clientX, e.touches[0].clientY]
                : [e.clientX, e.clientY];
            return { x: (cx - rect.left) * scaleX, y: (cy - rect.top) * scaleY };
        };

        const start = (e) => {
            e.preventDefault();
            state.isPainting = true;
            const c = getCoords(e);
            lastX = c.x; lastY = c.y;
            this.paintDot(c.x, c.y);
        };
        const move = (e) => {
            if (!state.isPainting) return;
            e.preventDefault();
            const c = getCoords(e);
            this.paintLine(lastX, lastY, c.x, c.y);
            lastX = c.x; lastY = c.y;
        };
        const end = () => {
            if (!state.isPainting) return;
            state.isPainting = false;
            this.saveMaskState();
        };

        canvas.addEventListener('mousedown', start);
        canvas.addEventListener('mousemove', move);
        canvas.addEventListener('mouseup', end);
        canvas.addEventListener('mouseleave', end);
        canvas.addEventListener('touchstart', start, { passive: false });
        canvas.addEventListener('touchmove', move, { passive: false });
        canvas.addEventListener('touchend', end);
    },

    paintDot(x, y) {
        const s = NAIState;
        const r = s.brushSize / 2;
        s.overlayCtx.fillStyle = 'rgba(255,255,255,0.25)';
        s.overlayCtx.beginPath(); s.overlayCtx.arc(x, y, r, 0, Math.PI * 2); s.overlayCtx.fill();
        s.maskCtx.fillStyle = '#ffffff';
        s.maskCtx.beginPath(); s.maskCtx.arc(x, y, r, 0, Math.PI * 2); s.maskCtx.fill();
    },

    paintLine(x1, y1, x2, y2) {
        const s = NAIState;
        const props = { lineWidth: s.brushSize, lineCap: 'round', lineJoin: 'round' };

        Object.assign(s.overlayCtx, props);
        s.overlayCtx.strokeStyle = 'rgba(255,255,255,0.25)';
        s.overlayCtx.beginPath(); s.overlayCtx.moveTo(x1, y1); s.overlayCtx.lineTo(x2, y2); s.overlayCtx.stroke();

        Object.assign(s.maskCtx, props);
        s.maskCtx.strokeStyle = '#ffffff';
        s.maskCtx.beginPath(); s.maskCtx.moveTo(x1, y1); s.maskCtx.lineTo(x2, y2); s.maskCtx.stroke();
    },

    saveMaskState() {
        const s = NAIState;
        if (!s.maskCanvas || !s.overlayCanvas) return;
        s.maskHistory = s.maskHistory.slice(0, s.maskHistoryIndex + 1);
        s.maskHistory.push({
            mask: s.maskCtx.getImageData(0, 0, s.maskCanvas.width, s.maskCanvas.height),
            overlay: s.overlayCtx.getImageData(0, 0, s.overlayCanvas.width, s.overlayCanvas.height),
        });
        s.maskHistoryIndex = s.maskHistory.length - 1;
        if (s.maskHistory.length > 30) { s.maskHistory.shift(); s.maskHistoryIndex--; }
    },

    undoMask() {
        const s = NAIState;
        if (s.maskHistoryIndex <= 0) return;
        s.maskHistoryIndex--;
        const snap = s.maskHistory[s.maskHistoryIndex];
        s.maskCtx.putImageData(snap.mask, 0, 0);
        s.overlayCtx.putImageData(snap.overlay, 0, 0);
    },

    clearMask() {
        const s = NAIState;
        if (!s.maskCanvas || !s.overlayCanvas) return;
        s.maskCtx.fillStyle = '#000000';
        s.maskCtx.fillRect(0, 0, s.maskCanvas.width, s.maskCanvas.height);
        s.overlayCtx.clearRect(0, 0, s.overlayCanvas.width, s.overlayCanvas.height);
        this.saveMaskState();
    },

    processMask(maskImage) {
        const width = maskImage.width;
        const height = maskImage.height;
        const smallWidth = Math.floor(width / 8);
        const smallHeight = Math.floor(height / 8);
    
        const smallCanvas = document.createElement('canvas');
        smallCanvas.width = smallWidth;
        smallCanvas.height = smallHeight;
        const smallCtx = smallCanvas.getContext('2d');
        smallCtx.imageSmoothingEnabled = false;
        smallCtx.drawImage(maskImage, 0, 0, width, height, 0, 0, smallWidth, smallHeight);

        const largeCanvas = document.createElement('canvas');
        largeCanvas.width = width;
        largeCanvas.height = height;
        const largeCtx = largeCanvas.getContext('2d');
        largeCtx.imageSmoothingEnabled = false;
        largeCtx.drawImage(smallCanvas, 0, 0, smallWidth, smallHeight, 0, 0, width, height);
    
        const imageData = largeCtx.getImageData(0, 0, width, height);
        const data = imageData.data;
        for (let i = 0; i < data.length; i += 4) {
            data[i] = data[i] < 255 ? 0 : 255;     // R
            data[i+1] = data[i+1] < 255 ? 0 : 255; // G
            data[i+2] = data[i+2] < 255 ? 0 : 255; // B
        }
        largeCtx.putImageData(imageData, 0, 0);
    
        return largeCanvas;
    },

    getMaskBase64() {
        if (!NAIState.maskCanvas) return null;
        const processed_mask = this.processMask(NAIState.maskCanvas);
        return NAIUtils.dataUrlToBase64(processed_mask.toDataURL('image/png'));
    },

    collectParams() {
        if (!NAIState.extraInputImage) return null;
        const mode = NAIState.extraInputMode;
        const result = { mode, image: NAIState.extraInputImage.b64 };

        if (mode === 'img2img') {
            const d = $('#extra-denoise');
            const n = $('#extra-noise');
            result.strength = d ? parseFloat(d.value) : 0.7;
            result.noise = n ? parseFloat(n.value) : 0;
        } else {
            const s = $('#extra-inpaint-strength');
            result.strength = s ? parseFloat(s.value) : 1;
            result.mask = this.getMaskBase64();
            const addOrig = $('#extra-add-original');
            result.add_original_image = addOrig ? addOrig.checked : true;
        }
        return result;
    },

    useLastGenerated() {
        const imgEl = document.querySelector('#result-image-area img');
        if (!imgEl || !imgEl.src) return;

        const dataUrl = imgEl.src;
        const b64 = NAIUtils.dataUrlToBase64(dataUrl);
        NAIState.extraInputImage = { file: null, dataUrl, b64 };

        // Reset mask when loading new image
        NAIState.maskHistory = [];
        NAIState.maskHistoryIndex = -1;
        NAIState.maskCanvas = null;
        NAIState.maskCtx = null;
        NAIState.overlayCanvas = null;
        NAIState.overlayCtx = null;

        this.renderPreview();
    },
    
    async handleUploadFromDataUrl(dataUrl) {
        const b64 = NAIUtils.dataUrlToBase64(dataUrl);
        NAIState.extraInputImage = { file: null, dataUrl, b64 };
        NAIState.maskHistory = [];
        NAIState.maskHistoryIndex = -1;
        NAIState.maskCanvas = null;
        NAIState.maskCtx = null;
        NAIState.overlayCanvas = null;
        NAIState.overlayCtx = null;
        this.renderPreview();
    },

    init() {
        // Toggle
        $('#extra-input-toggle').addEventListener('click', () => {
            const c = $('#extra-input-content');
            const a = $('#extra-input-arrow');
            const open = c.classList.toggle('open');
            a.textContent = open ? '▼' : '▶';
        });

        // Upload
        const input = $('#extra-input-file');
        const drop = $('#extra-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 = '';
        });

        // Mode
        $$('input[name="extra-mode"]').forEach(radio => {
            radio.addEventListener('change', () => this.setMode(radio.value));
        });

        // Clear
        $('#btn-extra-clear').addEventListener('click', () => this.clear());

        // Use last generated image
        $('#btn-extra-use-last').addEventListener('click', () => this.useLastGenerated());
    },
};