Greasy Fork

Greasy Fork is available in English.

公益酒馆ComfyUI插图脚本

移除直连模式,专注调度器并增加自定义缓存管理。

当前为 2025-09-07 提交的版本,查看 最新版本

您需要先安装一款用户脚本管理器扩展,例如 Tampermonkey 篡改猴Greasemonkey 油猴子Violentmonkey 暴力猴,才能安装此脚本。

您需要先安装一款用户脚本管理器扩展,例如 Tampermonkey 篡改猴,才能安装此脚本。

您需要先安装一款用户脚本管理器扩展,例如 Tampermonkey 篡改猴Violentmonkey 暴力猴,才能安装此脚本。

您需要先安装一款用户脚本管理器扩展,例如 Tampermonkey 篡改猴Userscripts ,才能安装此脚本。

您需要先安装一款用户脚本管理器扩展,例如 Tampermonkey 篡改猴,才能安装此脚本。

您需要先安装一款用户脚本管理器扩展后才能安装此脚本。

(我已经安装了用户脚本管理器,让我安装!)

您需要先安装一款用户样式管理器扩展,比如 Stylus,才能安装此样式。

您需要先安装一款用户样式管理器扩展,比如 Stylus,才能安装此样式。

您需要先安装一款用户样式管理器扩展,比如 Stylus,才能安装此样式。

您需要先安装一款用户样式管理器扩展后才能安装此样式。

您需要先安装一款用户样式管理器扩展后才能安装此样式。

您需要先安装一款用户样式管理器扩展后才能安装此样式。

(我已经安装了用户样式管理器,让我安装!)

// ==UserScript==
// @name         公益酒馆ComfyUI插图脚本
// @namespace    http://tampermonkey.net/
// @version      28.0 // 版本号递增,功能重构
// @license GPL
// @description  移除直连模式,专注调度器并增加自定义缓存管理。
// @author       feng zheng
// @match        *://*/*
// @grant        GM_addStyle
// @grant        GM_setValue
// @grant        GM_getValue
// @grant        GM_xmlhttpRequest
// @require      https://code.jquery.com/ui/1.13.2/jquery-ui.min.js
// ==/UserScript==


(function() {
    'use strict';

    // --- Configuration Constants ---
    const BUTTON_ID = 'comfyui-launcher-button';
    const PANEL_ID = 'comfyui-panel';
    const POLLING_TIMEOUT_MS = 120000;
    const POLLING_INTERVAL_MS = 3000;
    const STORAGE_KEY_IMAGES = 'comfyui_generated_images';
    const STORAGE_KEY_PROMPT_PREFIX = 'comfyui_prompt_prefix';
    const STORAGE_KEY_MAX_WIDTH = 'comfyui_image_max_width';
    const STORAGE_KEY_CACHE_LIMIT = 'comfyui_cache_limit'; // 新增:缓存上限的存储键
    const COOLDOWN_DURATION_MS = 60000;

    // --- Global Cooldown Variable ---
    let globalCooldownEndTime = 0;

    // --- Cached User Settings Variables ---
    let cachedSettings = {
        comfyuiUrl: '',
        startTag: 'image###',
        endTag: '###',
        promptPrefix: '',
        maxWidth: 600,
        cacheLimit: 20 // 新增:缓存上限
    };

    // --- Inject Custom CSS Styles ---
    GM_addStyle(`
        /* ... 您的所有 CSS 样式代码保持不变 ... */
        /* 新增:缓存状态显示样式 */
        #comfyui-cache-status {
            margin-top: 15px;
            margin-bottom: 10px;
            padding: 8px;
            background-color: rgba(0,0,0,0.2);
            border: 1px solid var(--SmartThemeBorderColor, #555);
            border-radius: 4px;
            text-align: center;
            font-size: 0.9em;
            color: #ccc;
        }
        /* 控制面板主容器样式 */
        #${PANEL_ID} {
            display: none; /* 默认隐藏 */
            position: fixed; /* 浮动窗口 */
            top: 50%;
            left: 50%;
            transform: translate(-50%, -50%); /* 默认居中显示 */
            width: 90vw; /* 移动设备上宽度 */
            max-width: 500px; /* 桌面设备上最大宽度 */
            z-index: 9999; /* 确保在顶层 */
            color: var(--SmartThemeBodyColor, #dcdcd2);
            background-color: var(--SmartThemeBlurTintColor, rgba(23, 23, 23, 0.9));
            border: 1px solid var(--SmartThemeBorderColor, rgba(0, 0, 0, 0.5));
            border-radius: 8px;
            box-shadow: 0 4px 15px var(--SmartThemeShadowColor, rgba(0, 0, 0, 0.5));
            padding: 15px;
            box-sizing: border-box;
            backdrop-filter: blur(var(--blurStrength, 10px));
            flex-direction: column;
        }

        /* 面板标题栏 */
        #${PANEL_ID} .panel-control-bar {
            /* 移除了 cursor: move; 使其不可拖动 */
            padding-bottom: 10px;
            margin-bottom: 15px;
            border-bottom: 1px solid var(--SmartThemeBorderColor, rgba(0, 0, 0, 0.5));
            display: flex;
            align-items: center;
            justify-content: space-between;
            flex-shrink: 0;
        }

        #${PANEL_ID} .panel-control-bar b { font-size: 1.2em; margin-left: 10px; }
        #${PANEL_ID} .floating_panel_close { cursor: pointer; font-size: 1.5em; }
        #${PANEL_ID} .floating_panel_close:hover { opacity: 0.7; }
        #${PANEL_ID} .comfyui-panel-content { overflow-y: auto; flex-grow: 1; padding-right: 5px; }

        /* 输入框和文本域样式 */
        #${PANEL_ID} input[type="text"],
        #${PANEL_ID} textarea,
        #${PANEL_ID} input[type="number"] { /* 包含数字输入框 */
            width: 100%;
            box-sizing: border-box;
            padding: 8px;
            border-radius: 4px;
            border: 1px solid var(--SmartThemeBorderColor, #555);
            background-color: rgba(0,0,0,0.2);
            color: var(--SmartThemeBodyColor, #dcdcd2);
            margin-bottom: 10px;
        }

        #${PANEL_ID} textarea { min-height: 150px; resize: vertical; }
        #${PANEL_ID} .workflow-info { font-size: 0.9em; color: #aaa; margin-top: -5px; margin-bottom: 10px;}

        /* 通用按钮样式 (用于测试连接和聊天内生成按钮) */
        .comfy-button {
            padding: 8px 12px;
            border: 1px solid black;
            border-radius: 4px;
            cursor: pointer;
            /* Modified: Changed button background to a gradient sky blue */
            background: linear-gradient(135deg, #87CEEB 0%, #00BFFF 100%); /* 天蓝色到深天蓝色渐变 */
            color: white;
            font-weight: 600;
            transition: opacity 0.3s, background 0.3s;
            flex-shrink: 0;
            font-size: 14px;
        }
        .comfy-button:disabled { opacity: 0.5; cursor: not-allowed; }
        .comfy-button:hover:not(:disabled) { opacity: 0.85; }

        /* 按钮状态样式 */
        .comfy-button.testing { background: #555; }
        .comfy-button.success { background: linear-gradient(135deg, #28a745 0%, #218838 100%); }
        .comfy-button.error   { background: linear-gradient(135deg, #dc3545 0%, #c82333 100%); }

        /* 特殊布局样式 */
        #comfyui-test-conn { position: relative; top: -5px; }
        .comfy-url-container { display: flex; gap: 10px; align-items: center; }
        .comfy-url-container input { flex-grow: 1; margin-bottom: 0; }
        #${PANEL_ID} label { display: block; margin-bottom: 5px; font-weight: bold; }
        #options > .options-content > a#${BUTTON_ID} { display: flex; align-items: center; gap: 10px; }

        /* 标记输入框容器样式 */
        #${PANEL_ID} .comfy-tags-container {
            display: flex;
            gap: 10px;
            align-items: flex-end;
            margin-top: 10px;
            margin-bottom: 10px;
        }
        #${PANEL_ID} .comfy-tags-container div { flex-grow: 1; }

        /* 聊天内按钮组容器 */
        .comfy-button-group {
            display: inline-flex;
            align-items: center;
            gap: 5px;
            margin: 5px 4px;
        }

        /* 生成的图片容器样式 */
        .comfy-image-container {
            margin-top: 10px;
            max-width: 100%; /* 默认允许图片最大宽度为容器的100% */
        }
        .comfy-image-container img {
            /* 注意:这里的max-width将由JavaScript直接设置,CSS变量作为备用或默认值 */
            max-width: var(--comfy-image-max-width, 100%);
            height: auto; /* 保持图片纵横比 */
            border-radius: 8px;
            border: 1px solid var(--SmartThemeBorderColor, #555);
        }

        /* 移动端适配 */
        @media (max-width: 1000px) {
            #${PANEL_ID} {
                top: 20px;
                left: 50%;
                transform: translateX(-50%);
                max-height: calc(100vh - 40px);
                width: 95vw;
            }
        }

        /* 定义一个CSS变量,用于动态控制图片最大宽度 */
        :root {
            --comfy-image-max-width: 600px; /* 默认图片最大宽度 */
        }
    `);

    // A flag to prevent duplicate execution from touchstart and click
    let lastTapTimestamp = 0;
    const TAP_THRESHOLD = 300; // milliseconds to prevent double taps/clicks

    function createComfyUIPanel() {
        if (document.getElementById(PANEL_ID)) return;
        // *** UI 修改:移除工作流文本域,替换为缓存管理 ***
        const panelHTML = `
            <div id="${PANEL_ID}">
                <div class="panel-control-bar">
                    <i class="fa-fw fa-solid fa-grip drag-grabber"></i>
                    <b>ComfyUI 生成设置</b>
                    <i class="fa-fw fa-solid fa-circle-xmark floating_panel_close"></i>
                </div>
                <div class="comfyui-panel-content">
                    <label for="comfyui-url">调度器 URL</label>
                    <div class="comfy-url-container">
                        <input id="comfyui-url" type="text" placeholder="例如: http://127.0.0.1:5001">
                        <button id="comfyui-test-conn" class="comfy-button">测试连接</button>
                    </div>
                    <div class="comfy-tags-container">
                        <div>
                            <label for="comfyui-start-tag">开始标记</label>
                            <input id="comfyui-start-tag" type="text">
                        </div>
                        <div>
                            <label for="comfyui-end-tag">结束标记</label>
                            <input id="comfyui-end-tag" type="text">
                        </div>
                    </div>
                    <div>
                        <label for="comfyui-prompt-prefix">提示词固定前缀 (LoRA等):</label>
                        <input id="comfyui-prompt-prefix" type="text" placeholder="例如: <lora:cool_style:0.8> ">
                    </div>
                    <div>
                        <label for="comfyui-max-width">最大图片宽度 (px):</label>
                        <input id="comfyui-max-width" type="number" placeholder="例如: 600" min="100">
                    </div>
                    <div>
                        <label for="comfyui-cache-limit">最大缓存数量:</label>
                        <input id="comfyui-cache-limit" type="number" placeholder="例如: 20" min="1" max="100">
                    </div>
                    <div id="comfyui-cache-status">当前缓存: ...</div>
                    <button id="comfyui-clear-cache" class="comfy-button error" style="margin-top: 15px; width: 100%;">删除所有图片缓存</button>
                </div>
            </div>
        `;
        document.body.insertAdjacentHTML('beforeend', panelHTML);
        initPanelLogic();
    }

    async function updateCacheStatusDisplay() {
        const display = document.getElementById('comfyui-cache-status');
        if (!display) return;
        const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
        const count = Object.keys(records).length;
        display.textContent = `当前缓存: ${count} / ${cachedSettings.cacheLimit} 张`;
    }

    function initPanelLogic() {
        const panel = document.getElementById(PANEL_ID);
        const closeButton = panel.querySelector('.floating_panel_close');
        const testButton = document.getElementById('comfyui-test-conn');
        const clearCacheButton = document.getElementById('comfyui-clear-cache');
        const urlInput = document.getElementById('comfyui-url');
        const startTagInput = document.getElementById('comfyui-start-tag');
        const endTagInput = document.getElementById('comfyui-end-tag');
        const promptPrefixInput = document.getElementById('comfyui-prompt-prefix');
        const maxWidthInput = document.getElementById('comfyui-max-width');
        const cacheLimitInput = document.getElementById('comfyui-cache-limit'); // 新增

        closeButton.addEventListener('click', () => { panel.style.display = 'none'; });

        testButton.addEventListener('click', () => {
            let url = urlInput.value.trim();
            if (!url) {
                if (typeof toastr !== 'undefined') toastr.warning('请输入调度器的URL。');
                return;
            }
            if (!url.startsWith('http://') && !url.startsWith('https://')) { url = 'http://' + url; }
            if (url.endsWith('/')) { url = url.slice(0, -1); }
            urlInput.value = url;
            const testUrl = url + '/system_stats';
            if (typeof toastr !== 'undefined') toastr.info('正在尝试连接服务...');
            testButton.className = 'comfy-button testing';
            testButton.disabled = true;
            GM_xmlhttpRequest({
                method: "GET", url: testUrl, timeout: 5000,
                onload: (res) => {
                    testButton.disabled = false;
                    testButton.className = res.status === 200 ? 'comfy-button success' : 'comfy-button error';
                    if (typeof toastr !== 'undefined') {
                        if(res.status === 200) toastr.success('连接成功!');
                        else toastr.error(`连接失败!状态: ${res.status}`);
                    }
                },
                onerror: () => {
                    testButton.disabled = false; testButton.className = 'comfy-button error';
                    if (typeof toastr !== 'undefined') toastr.error('连接错误!');
                },
                ontimeout: () => {
                    testButton.disabled = false; testButton.className = 'comfy-button error';
                    if (typeof toastr !== 'undefined') toastr.error('连接超时!');
                }
            });
        });

        clearCacheButton.addEventListener('click', async () => {
            if (confirm('您确定要删除所有已生成的图片缓存吗?')) {
                await GM_setValue(STORAGE_KEY_IMAGES, {});
                await updateCacheStatusDisplay(); // 更新显示
                if (typeof toastr !== 'undefined') toastr.success('图片缓存已清空!');
            }
        });

        loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput).then(() => {
            applyCurrentMaxWidthToAllImages();
        });

        [urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput].forEach(input => {
            input.addEventListener('input', async () => {
                if (input === urlInput) testButton.className = 'comfy-button';
                await saveSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput);
                if (input === maxWidthInput) applyCurrentMaxWidthToAllImages();
            });
        });
    }

    async function loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput) {
        cachedSettings.comfyuiUrl = await GM_getValue('comfyui_url', 'http://127.0.0.1:5001');
        cachedSettings.startTag = await GM_getValue('comfyui_start_tag', 'image###');
        cachedSettings.endTag = await GM_getValue('comfyui_end_tag', '###');
        cachedSettings.promptPrefix = await GM_getValue(STORAGE_KEY_PROMPT_PREFIX, '');
        cachedSettings.maxWidth = await GM_getValue(STORAGE_KEY_MAX_WIDTH, 600);
        cachedSettings.cacheLimit = await GM_getValue(STORAGE_KEY_CACHE_LIMIT, 20); // 新增

        urlInput.value = cachedSettings.comfyuiUrl;
        startTagInput.value = cachedSettings.startTag;
        endTagInput.value = cachedSettings.endTag;
        promptPrefixInput.value = cachedSettings.promptPrefix;
        maxWidthInput.value = cachedSettings.maxWidth;
        cacheLimitInput.value = cachedSettings.cacheLimit; // 新增

        document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
        await updateCacheStatusDisplay(); // 加载设置后更新显示
    }

    async function saveSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput) {
        cachedSettings.comfyuiUrl = urlInput.value.trim();
        cachedSettings.startTag = startTagInput.value;
        cachedSettings.endTag = endTagInput.value;
        cachedSettings.promptPrefix = promptPrefixInput.value.trim();
        const newMaxWidth = parseInt(maxWidthInput.value);
        cachedSettings.maxWidth = isNaN(newMaxWidth) ? 600 : newMaxWidth;
        const newCacheLimit = parseInt(cacheLimitInput.value); // 新增
        cachedSettings.cacheLimit = isNaN(newCacheLimit) ? 20 : newCacheLimit; // 新增

        await GM_setValue('comfyui_url', cachedSettings.comfyuiUrl);
        await GM_setValue('comfyui_start_tag', cachedSettings.startTag);
        await GM_setValue('comfyui_end_tag', cachedSettings.endTag);
        await GM_setValue(STORAGE_KEY_PROMPT_PREFIX, cachedSettings.promptPrefix);
        await GM_setValue(STORAGE_KEY_MAX_WIDTH, cachedSettings.maxWidth);
        await GM_setValue(STORAGE_KEY_CACHE_LIMIT, cachedSettings.cacheLimit); // 新增

        document.documentElement.style.setProperty('--comfy-image-max-width', cachedSettings.maxWidth + 'px');
        await updateCacheStatusDisplay(); // 保存设置后更新显示
    }

    async function applyCurrentMaxWidthToAllImages() {
        const images = document.querySelectorAll('.comfy-image-container img');
        const maxWidthPx = (cachedSettings.maxWidth || 600) + 'px';
        images.forEach(img => { img.style.maxWidth = maxWidthPx; });
    }

    function addMainButton() {
        if (document.getElementById(BUTTON_ID)) return;
        const optionsMenuContent = document.querySelector('#options .options-content');
        if (optionsMenuContent) {
            const continueButton = optionsMenuContent.querySelector('#option_continue');
            if (continueButton) {
                const comfyButton = document.createElement('a');
                comfyButton.id = BUTTON_ID;
                comfyButton.className = 'interactable';
                comfyButton.innerHTML = `<i class="fa-lg fa-solid fa-image"></i><span>ComfyUI生图</span>`;
                comfyButton.style.cursor = 'pointer';
                comfyButton.addEventListener('click', (event) => {
                    event.preventDefault();
                    document.getElementById(PANEL_ID).style.display = 'flex';
                    document.getElementById('options').style.display = 'none';
                });
                continueButton.parentNode.insertBefore(comfyButton, continueButton.nextSibling);
            }
        }
    }

    // --- Helper and Cache Management Functions ---
    // ... (Your existing helper functions: escapeRegex, generateClientId, simpleHash)
    function escapeRegex(string) {
        return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
    }

    function generateClientId() {
        return 'client-' + Math.random().toString(36).substring(2, 15) + Math.random().toString(36).substring(2, 15);
    }

    function simpleHash(str) {
        let hash = 0;
        for (let i = 0; i < str.length; i++) {
            const char = str.charCodeAt(i);
            hash = (hash << 5) - hash + char;
            hash |= 0;
        }
        return 'comfy-id-' + Math.abs(hash).toString(36);
    }


    function fetchImageAsBase64(imageUrl) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'GET', url: imageUrl, responseType: 'blob', timeout: 30000,
                onload: (response) => {
                    if (response.status === 200) {
                        const reader = new FileReader();
                        reader.onloadend = () => resolve(reader.result);
                        reader.onerror = (err) => reject(new Error('FileReader error: ' + err));
                        reader.readAsDataURL(response.response);
                    } else {
                        reject(new Error(`获取图片失败,状态: ${response.status}`));
                    }
                },
                onerror: (err) => reject(new Error('网络错误,无法下载图片: ' + err)),
                ontimeout: () => reject(new Error('下载图片超时。'))
            });
        });
    }

    async function saveImageRecord(generationId, imageBase64Data) {
        let records = await GM_getValue(STORAGE_KEY_IMAGES, {});
        if (records.hasOwnProperty(generationId)) {
            delete records[generationId];
        }
        records[generationId] = imageBase64Data;

        const keys = Object.keys(records);
        // *** 使用用户自定义的缓存上限 ***
        if (keys.length > cachedSettings.cacheLimit) {
            const keysToDelete = keys.slice(0, keys.length - cachedSettings.cacheLimit);
            keysToDelete.forEach(key => { delete records[key]; });
            console.log(`缓存已满,删除了 ${keysToDelete.length} 条旧记录。`);
            if (typeof toastr !== 'undefined') toastr.info(`缓存已更新,旧记录已清理。`);
        }
        await GM_setValue(STORAGE_KEY_IMAGES, records);
        await updateCacheStatusDisplay(); // 更新显示
    }

    async function deleteImageRecord(generationId) {
        const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
        delete records[generationId];
        await GM_setValue(STORAGE_KEY_IMAGES, records);
        await updateCacheStatusDisplay(); // 更新显示
    }

    // --- Chat Message Processing and Image Generation ---
    // ... (Your existing chat processing functions: handleComfyButtonClick, processMessageForComfyButton, etc.)
    function handleComfyButtonClick(event, isTouch = false) {
        const button = event.target.closest('.comfy-chat-generate-button');
        if (!button) return;

        if (isTouch) {
            event.preventDefault();
            const now = Date.now();
            if (now - lastTapTimestamp < TAP_THRESHOLD) return;
            lastTapTimestamp = now;
            onGenerateButtonClickLogic(button);
        } else {
            if (Date.now() - lastTapTimestamp < TAP_THRESHOLD) return;
            onGenerateButtonClickLogic(button);
        }
    }

    async function processMessageForComfyButton(messageNode, savedImagesCache) {
        const mesText = messageNode.querySelector('.mes_text');
        if (!mesText) return;

        const startTag = cachedSettings.startTag;
        const endTag = cachedSettings.endTag;
        if (!startTag || !endTag) return;

        const escapedStartTag = escapeRegex(startTag);
        const escapedEndTag = escapeRegex(endTag);
        const regex = new RegExp(escapedStartTag + '([\\s\\S]*?)' + escapedEndTag, 'g');
        const currentHtml = mesText.innerHTML;

        if (regex.test(currentHtml) && !mesText.querySelector('.comfy-button-group')) {
            mesText.innerHTML = currentHtml.replace(regex, (match, prompt) => {
                const cleanPrompt = prompt.trim();
                const encodedPrompt = cleanPrompt.replace(/"/g, '"');
                const generationId = simpleHash(cleanPrompt);
                return `<span class="comfy-button-group" data-generation-id="${generationId}">
                            <button class="comfy-button comfy-chat-generate-button" data-prompt="${encodedPrompt}">开始生成</button>
                        </span>`;
            });
        }

        const buttonGroups = mesText.querySelectorAll('.comfy-button-group');
        buttonGroups.forEach(group => {
            if (group.dataset.listenerAttached) return;

            const generationId = group.dataset.generationId;
            const generateButton = group.querySelector('.comfy-chat-generate-button');

            if (Date.now() < globalCooldownEndTime) {
                generateButton.dataset.cooldownEnd = globalCooldownEndTime.toString();
                startCooldownCountdown(generateButton, globalCooldownEndTime);
            } else if (savedImagesCache[generationId]) {
                displayImage(group, savedImagesCache[generationId]);
                setupGeneratedState(generateButton, generationId);
            }
            group.dataset.listenerAttached = 'true';
        });
    }

    function setupGeneratedState(generateButton, generationId) {
        generateButton.textContent = '重新生成';
        generateButton.disabled = false;
        generateButton.classList.remove('testing', 'success', 'error');
        delete generateButton.dataset.cooldownEnd;

        const group = generateButton.closest('.comfy-button-group');
        let deleteButton = group.querySelector('.comfy-delete-button');
        if (!deleteButton) {
            deleteButton = document.createElement('button');
            deleteButton.textContent = '删除';
            deleteButton.className = 'comfy-button error comfy-delete-button';
            deleteButton.addEventListener('click', async () => {
                await deleteImageRecord(generationId);
                const imageContainer = group.nextElementSibling;
                if (imageContainer && imageContainer.classList.contains('comfy-image-container')) {
                    imageContainer.remove();
                }
                deleteButton.remove();
                generateButton.textContent = '开始生成';
            });
            generateButton.insertAdjacentElement('afterend', deleteButton);
        }
    }


    async function onGenerateButtonClickLogic(button) {
        const group = button.closest('.comfy-button-group');
        let prompt = button.dataset.prompt;
        const generationId = group.dataset.generationId;

        if (button.disabled) return;
        if (Date.now() < globalCooldownEndTime) {
            const remainingTime = Math.ceil((globalCooldownEndTime - Date.now()) / 1000);
            if (typeof toastr !== 'undefined') toastr.warning(`请稍候,冷却中 (${remainingTime}s)。`);
            return;
        }

        button.textContent = '生成中...';
        button.disabled = true;
        button.classList.add('testing');
        const deleteButton = group.querySelector('.comfy-delete-button');
        if (deleteButton) deleteButton.style.display = 'none';
        const oldImageContainer = group.nextElementSibling;

        try {
            const url = cachedSettings.comfyuiUrl;
            const promptPrefix = cachedSettings.promptPrefix;
            if (!url) throw new Error('调度器 URL 未配置。');
            if (promptPrefix) prompt = promptPrefix + ' ' + prompt;

            const clientId = generateClientId();

            // *** 逻辑简化:总是使用调度器模式 ***
            if (typeof toastr !== 'undefined') toastr.info('正在向调度器发送请求...');
            const promptResponse = await sendPromptRequestToScheduler(url, {
                client_id: clientId,
                positive_prompt: prompt
            });
            if (promptResponse.assigned_instance_name) {
                if (typeof toastr !== 'undefined') toastr.success(`任务已分配到: ${promptResponse.assigned_instance_name} (队列: ${promptResponse.assigned_instance_queue_size})`);
            }

            const promptId = promptResponse.prompt_id;
            if (!promptId) throw new Error('调度器未返回有效的任务 ID。');

            const finalHistory = await pollForResult(url, promptId);
            const imageUrl = findImageUrlInHistory(finalHistory, promptId, url);
            if (!imageUrl) throw new Error('未在结果中找到图片 URL。');

            if (typeof toastr !== 'undefined') toastr.info('正在获取图片数据并缓存...');
            const imageBase64Data = await fetchImageAsBase64(imageUrl);
            if (!imageBase64Data) throw new Error('无法获取图片数据。');

            if (oldImageContainer) oldImageContainer.remove();
            displayImage(group, imageBase64Data);
            await saveImageRecord(generationId, imageBase64Data);

            button.textContent = '生成成功';
            button.classList.remove('testing', 'success');
            setTimeout(() => {
                setupGeneratedState(button, generationId);
                if (deleteButton) deleteButton.style.display = 'inline-flex';
            }, 2000);

        } catch (e) {
            console.error('ComfyUI 生图脚本错误:', e);
            let displayMessage = '图片生成失败,请检查服务。';
            let isRateLimitError = false;
            let actualCooldownSeconds = COOLDOWN_DURATION_MS / 1000;

            const rateLimitMatch = e.message.match(/请在 (\d+) 秒后重试。/);
            if (rateLimitMatch && rateLimitMatch[1]) {
                actualCooldownSeconds = parseInt(rateLimitMatch[1], 10);
                displayMessage = `请求频率过高,请在 ${actualCooldownSeconds} 秒后重试。`;
                isRateLimitError = true;
            } else if (e.message.includes('请求频率过高')) {
                displayMessage = '请求频率过高,请稍后再试。';
                isRateLimitError = true;
            } else if (e.message.includes('轮询结果超时')) {
                displayMessage = '生成任务超时。';
            } else {
                const backendErrorMatch = e.message.match(/error:\s*"(.*?)"/);
                if (backendErrorMatch && backendErrorMatch[1]) {
                    displayMessage = `调度器错误: ${backendErrorMatch[1]}`;
                }
            }

            if (typeof toastr !== 'undefined') toastr.error(displayMessage);
            if (deleteButton) deleteButton.style.display = 'inline-flex';

            if (isRateLimitError) {
                const newCooldownEndTime = Date.now() + (actualCooldownSeconds * 1000);
                globalCooldownEndTime = newCooldownEndTime;
                applyGlobalCooldown(newCooldownEndTime);
            } else {
                button.textContent = '生成失败';
                button.classList.add('error');
                setTimeout(() => {
                    const wasRegenerating = !!group.querySelector('.comfy-delete-button');
                    button.classList.remove('error');
                    if (wasRegenerating) {
                        setupGeneratedState(button, generationId);
                    } else {
                        button.textContent = '开始生成';
                        button.disabled = false;
                    }
                }, 3000);
            }
        }
    }

    // ... (Your other functions like applyGlobalCooldown, startCooldownCountdown, sendPromptRequestToScheduler, pollForResult, findImageUrlInHistory, displayImage remain the same)
    function applyGlobalCooldown(endTime) {
        const allGenerateButtons = document.querySelectorAll('.comfy-chat-generate-button');
        allGenerateButtons.forEach(button => {
            button.dataset.cooldownEnd = endTime.toString();
            startCooldownCountdown(button, endTime);
        });
    }

    function startCooldownCountdown(button, endTime) {
        button.disabled = true;
        button.classList.remove('success', 'error', 'testing');
        const updateCountdown = () => {
            const remainingTime = Math.max(0, endTime - Date.now());
            const seconds = Math.ceil(remainingTime / 1000);
            if (seconds > 0) {
                button.textContent = `冷却中 (${seconds}s)`;
                setTimeout(updateCountdown, 1000);
            } else {
                button.disabled = false;
                delete button.dataset.cooldownEnd;
                const group = button.closest('.comfy-button-group');
                const generationId = group.dataset.generationId;
                const deleteButtonPresent = group.querySelector('.comfy-delete-button');
                if (deleteButtonPresent) {
                    setupGeneratedState(button, generationId);
                } else {
                    button.textContent = '开始生成';
                }
            }
        };
        updateCountdown();
    }

    // --- API Request Functions (sendPromptRequestToScheduler, sendPromptRequestDirect, etc.) ---
    // ... (Your existing API functions remain largely unchanged)
    function sendPromptRequestToScheduler(url, payload) {
        return new Promise((resolve, reject) => {
            GM_xmlhttpRequest({
                method: 'POST',
                url: `${url}/generate`,
                headers: { 'Content-Type': 'application/json' },
                data: JSON.stringify(payload),
                timeout: 10000, // 从 30s 减少到 10s
                onload: (res) => {
                    // *** 关键修改: 处理 202 Accepted 状态码 ***
                    if (res.status === 202) {
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至调度器,任务已在后台排队。');
                        let responseData = {};
                        try {
                             responseData = JSON.parse(res.responseText);
                        } catch (e) {
                             console.warn('调度器 202 响应不是有效的 JSON 或为空。继续使用空响应数据。', e);
                        }
                        // 调度器 202 响应中现在应该包含 prompt_id, assigned_instance_name, assigned_instance_queue_size
                        resolve({
                            prompt_id: responseData.prompt_id,
                            message: responseData.message,
                            assigned_instance_name: responseData.assigned_instance_name, // 新增
                            assigned_instance_queue_size: responseData.assigned_instance_queue_size // 新增
                        });
                    } else if (res.status === 200) { // 兼容旧版调度器或同步返回 prompt_id 的情况
                        if (typeof toastr !== 'undefined') toastr.info('请求已发送至调度器,排队中...');
                        resolve(JSON.parse(res.responseText));
                    }
                    else {
                        let errorMessage = '';
                        try {
                            const errorJson = JSON.parse(res.responseText);
                            if (errorJson && errorJson.error) {
                                // 如果是JSON错误,直接使用其error字段
                                errorMessage = errorJson.error;
                            } else {
                                // 否则,使用状态码和原始响应文本
                                errorMessage = `调度器 API 错误: ${res.statusText || res.status} - ${res.responseText}`;
                            }
                        } catch (parseError) {
                            // 如果响应文本不是JSON,直接作为错误信息
                            errorMessage = `调度器 API 错误: ${res.statusText || res.status} - ${res.responseText}`;
                        }
                        reject(new Error(errorMessage));
                    }
                },
                onerror: (e) => reject(new Error('无法连接到调度器 API。请检查URL和网络连接。详细错误:' + (e.responseText || e.statusText || e.status))),
                ontimeout: () => reject(new Error('连接调度器 API 超时。请检查网络。')),
            });
        });
    }


    function pollForResult(url, promptId) {
        return new Promise((resolve, reject) => {
            const startTime = Date.now();
            const poller = setInterval(() => {
                if (Date.now() - startTime > POLLING_TIMEOUT_MS) {
                    clearInterval(poller);
                    // 隐藏敏感信息
                    reject(new Error('轮询结果超时。任务可能仍在处理中或已失败。请查看调度器日志了解更多信息。'));
                    return;
                }
                GM_xmlhttpRequest({
                    method: 'GET',
                    url: `${url}/history/${promptId}`,
                    timeout: 65000, // 显式设置超时时间,略大于调度器的代理超时
                    onload: (res) => {
                        if (res.status === 200) {
                            const history = JSON.parse(res.responseText);
                            // 检查历史记录中是否存在该 promptId 并且其 outputs 不为空
                            if (history[promptId] && Object.keys(history[promptId].outputs).length > 0) {
                                clearInterval(poller);
                                resolve(history);
                            } else {
                                // 即使 200,如果 outputs 为空,也可能意味着任务仍在进行中
                                console.info(`轮询历史记录 ${promptId}: 任务仍在进行中。`); // 使用 console.info 避免频繁弹窗
                            }
                        } else if (res.status === 404) {
                             clearInterval(poller);
                             // 隐藏敏感信息
                             reject(new Error(`轮询结果失败: 任务 ID ${promptId} 未找到或已过期。请查看调度器日志了解更多信息。`));
                        }
                        else {
                            clearInterval(poller);
                            // 隐藏敏感信息,只显示状态码和通用信息
                            reject(new Error(`轮询结果失败: 后端服务返回状态码 ${res.status}。请查看调度器日志了解更多信息。`));
                        }
                    },
                    onerror: (e) => {
                        clearInterval(poller);
                        // 隐藏敏感信息,提供通用网络错误提示
                        reject(new Error('轮询结果网络错误或调度器/ComfyUI无响应。请检查网络连接或调度器日志。详细错误:' + (e.responseText || e.statusText || e.status)));
                    },
                    ontimeout: () => { // 处理单个轮询请求的超时
                        clearInterval(poller);
                        // 隐藏敏感信息
                        reject(new Error(`单个轮询请求超时。调度器在历史记录接口处无响应。请检查调度器日志了解更多信息。`));
                    }
                });
            }, POLLING_INTERVAL_MS);
        });
    }

    function findImageUrlInHistory(history, promptId, baseUrl) {
        const outputs = history[promptId]?.outputs;
        if (!outputs) return null;

        for (const nodeId in outputs) {
            if (outputs.hasOwnProperty(nodeId) && outputs[nodeId].images) {
                const image = outputs[nodeId].images[0];
                if (image) {
                    const params = new URLSearchParams({
                        filename: image.filename,
                        subfolder: image.subfolder,
                        type: image.type,
                        prompt_id: promptId // 传递 prompt_id 给 /view 路由
                    });
                    return `${baseUrl}/view?${params.toString()}`;
                }
            }
        }
        return null;
    }

    async function displayImage(anchorElement, imageBase64Data) {
        let container = anchorElement.nextElementSibling;
        if (!container || !container.classList.contains('comfy-image-container')) {
            container = document.createElement('div');
            container.className = 'comfy-image-container';
            const img = document.createElement('img');
            img.alt = 'ComfyUI 生成的图片';
            container.appendChild(img);
            anchorElement.insertAdjacentElement('afterend', container);
        }
        const imgElement = container.querySelector('img');
        imgElement.src = imageBase64Data;
        imgElement.style.maxWidth = (cachedSettings.maxWidth || 600) + 'px';
    }


    // --- Main Execution Logic ---
    // ... (Main logic functions: createComfyUIPanel, chatObserver, observeChat, etc.)
    // ... (Your existing API functions remain largely unchanged)
    // --- Main Execution Logic ---

    createComfyUIPanel();

    const chatObserver = new MutationObserver(async (mutations) => {
        const nodesToProcess = new Set();
        for (const mutation of mutations) {
            mutation.addedNodes.forEach(node => {
                if (node.nodeType === Node.ELEMENT_NODE) {
                    if (node.matches('.mes')) nodesToProcess.add(node);
                    node.querySelectorAll('.mes').forEach(mes => nodesToProcess.add(mes));
                }
            });
            if (mutation.target.nodeType === Node.ELEMENT_NODE && mutation.target.closest('.mes')) {
                 nodesToProcess.add(mutation.target.closest('.mes'));
            }
        }

        if (nodesToProcess.size > 0) {
            const savedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});
            await loadSettingsFromStorageAndApplyToCache();
            nodesToProcess.forEach(node => {
                const mesTextElement = node.querySelector('.mes_text');
                if (mesTextElement && !mesTextElement.dataset.listenersAttached) {
                    mesTextElement.addEventListener('touchstart', (event) => handleComfyButtonClick(event, true), { passive: false });
                    mesTextElement.addEventListener('click', (event) => handleComfyButtonClick(event, false));
                    mesTextElement.dataset.listenersAttached = 'true';
                }
                processMessageForComfyButton(node, savedImages);
            });
        }
    });

    async function loadSettingsFromStorageAndApplyToCache() {
        cachedSettings.comfyuiUrl = await GM_getValue('comfyui_url', 'http://127.0.0.1:5001');
        cachedSettings.startTag = await GM_getValue('comfyui_start_tag', 'image###');
        cachedSettings.endTag = await GM_getValue('comfyui_end_tag', '###');
        cachedSettings.promptPrefix = await GM_getValue(STORAGE_KEY_PROMPT_PREFIX, '');
        cachedSettings.maxWidth = await GM_getValue(STORAGE_KEY_MAX_WIDTH, 600);
        cachedSettings.cacheLimit = await GM_getValue(STORAGE_KEY_CACHE_LIMIT, 20); // 新增
        document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
    }

    function observeChat() {
        const chatElement = document.getElementById('chat');
        if (chatElement) {
            loadSettingsFromStorageAndApplyToCache().then(async () => {
                const initialSavedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});
                chatElement.querySelectorAll('.mes').forEach(node => {
                    const mesTextElement = node.querySelector('.mes_text');
                    if (mesTextElement && !mesTextElement.dataset.listenersAttached) {
                        mesTextElement.addEventListener('touchstart', (event) => handleComfyButtonClick(event, true), { passive: false });
                        mesTextElement.addEventListener('click', (event) => handleComfyButtonClick(event, false));
                        mesTextElement.dataset.listenersAttached = 'true';
                    }
                    processMessageForComfyButton(node, initialSavedImages);
                });
                chatObserver.observe(chatElement, { childList: true, subtree: true });
            });
        } else {
            setTimeout(observeChat, 500);
        }
    }

    const optionsObserver = new MutationObserver(() => {
        const optionsMenu = document.getElementById('options');
        if (optionsMenu && optionsMenu.style.display !== 'none') {
            addMainButton();
        }
    });

    window.addEventListener('load', () => {
        observeChat();
        const body = document.querySelector('body');
        if (body) {
            optionsObserver.observe(body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style'] });
        }
    });

})();