移除轮询，使用WebSocket实时接收生成结果，并增加API密钥认证、模型路由和默认模型选择功能。
// ==UserScript==
// @name 公益酒馆ComfyUI插图脚本 (WebSocket实时版)
// @namespace http://tampermonkey.net/
// @version 31.0 // 版本号递增，新增API密钥、模型选择和默认模型设置
// @license GPL
// @description 移除轮询，使用WebSocket实时接收生成结果，并增加API密钥认证、模型路由和默认模型选择功能。
// @author feng zheng (升级 by Gemini)
// @match *://*/*
// @grant GM_addStyle
// @grant GM_setValue
// @grant GM_getValue
// @grant GM_xmlhttpRequest
// @require https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.7.5/socket.io.min.js
// @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 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 State Variables ---
let globalCooldownEndTime = 0;
let socket = null;
let activePrompts = {}; // 存储 prompt_id -> { button, generationId } 的映射
// --- Cached User Settings ---
let cachedSettings = {
comfyuiUrl: '',
startTag: 'image###',
endTag: '###',
promptPrefix: '',
maxWidth: 600,
cacheLimit: 20,
apiKey: '', // 新增
defaultModel: '' // 新增
};
// --- 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 { 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"], #${PANEL_ID} select { 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; 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%; }
.comfy-image-container img { 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; }
`);
// --- WebSocket & State Management ---
function connectWebSocket() {
if (socket && socket.connected) return;
const schedulerUrl = new URL(cachedSettings.comfyuiUrl);
const wsUrl = `${schedulerUrl.protocol}//${schedulerUrl.host}`;
if (typeof io === 'undefined') {
console.error('Socket.IO 客户端库未加载！');
if (typeof toastr !== 'undefined') toastr.error('实时通信库加载失败！');
return;
}
socket = io(wsUrl, { reconnectionAttempts: 5, timeout: 20000 });
socket.on('connect', () => {
console.log('成功连接到调度器 WebSocket！');
if (typeof toastr !== 'undefined') toastr.success('已建立实时连接！');
});
socket.on('disconnect', () => {
console.log('与调度器 WebSocket 断开连接。');
});
socket.on('generation_complete', (data) => {
const { prompt_id, status, imageUrl, error } = data;
const promptInfo = activePrompts[prompt_id];
if (!promptInfo) return;
const { button, generationId } = promptInfo;
const group = button.closest('.comfy-button-group');
if (status === 'success' && imageUrl) {
if (typeof toastr !== 'undefined') toastr.info(`图片已生成！`);
displayImage(group, imageUrl);
cacheImageInBackground(generationId, imageUrl);
button.textContent = '生成成功';
button.classList.remove('testing');
button.classList.add('success');
setTimeout(() => {
setupGeneratedState(button, generationId);
}, 2000);
} else {
if (typeof toastr !== 'undefined') toastr.error(`生成失败: ${error || '未知错误'}`);
button.textContent = '生成失败';
button.classList.remove('testing');
button.classList.add('error');
setTimeout(() => {
button.disabled = false;
button.classList.remove('error');
button.textContent = group.querySelector('.comfy-delete-button') ? '重新生成' : '开始生成';
}, 3000);
}
delete activePrompts[prompt_id];
});
}
// --- Core Application Logic (UI, Settings, Image Handling) ---
// A flag to prevent duplicate execution
let lastTapTimestamp = 0;
const TAP_THRESHOLD = 300;
function createComfyUIPanel() {
if (document.getElementById(PANEL_ID)) return;
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-api-key">API 密钥:</label><input id="comfyui-api-key" type="text" placeholder="在此输入您的密钥"></div>
<div>
<label for="comfyui-default-model">默认模型 (不指定时生效):</label>
<select id="comfyui-default-model">
<option value="">自动选择</option>
<option value="waiNSFWIllustrious_v140">waiNSFWIllustrious_v140</option>
<option value="Pony_alpha">Pony_alpha</option>
</select>
</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');
const apiKeyInput = document.getElementById('comfyui-api-key');
const defaultModelSelect = document.getElementById('comfyui-default-model');
closeButton.addEventListener('click', () => { panel.style.display = 'none'; });
testButton.addEventListener('click', () => {
let url = urlInput.value.trim();
if (!url) { toastr?.warning('请输入调度器的URL。'); return; }
if (!url.startsWith('http')) { url = 'http://' + url; }
if (url.endsWith('/')) { url = url.slice(0, -1); }
urlInput.value = url;
const testUrl = url + '/system_stats';
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 (res.status === 200) toastr?.success('连接成功！');
else toastr?.error(`连接失败！状态: ${res.status}`);
},
onerror: () => {
testButton.disabled = false; testButton.className = 'comfy-button error';
toastr?.error('连接错误！');
},
ontimeout: () => {
testButton.disabled = false; testButton.className = 'comfy-button error';
toastr?.error('连接超时！');
}
});
});
clearCacheButton.addEventListener('click', async () => {
if (confirm('您确定要删除所有已生成的图片缓存吗？')) {
await GM_setValue(STORAGE_KEY_IMAGES, {});
await updateCacheStatusDisplay();
toastr?.success('图片缓存已清空！');
}
});
loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect).then(() => {
applyCurrentMaxWidthToAllImages();
});
[urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect].forEach(input => {
const eventType = input.tagName.toLowerCase() === 'select' ? 'change' : 'input';
input.addEventListener(eventType, async () => {
if (input === urlInput) testButton.className = 'comfy-button';
await saveSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect);
if (input === maxWidthInput) applyCurrentMaxWidthToAllImages();
if (input === urlInput) {
if (socket) socket.disconnect();
connectWebSocket();
}
});
});
}
async function loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect) {
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);
cachedSettings.apiKey = await GM_getValue('comfyui_api_key', '');
cachedSettings.defaultModel = await GM_getValue('comfyui_default_model', '');
urlInput.value = cachedSettings.comfyuiUrl;
startTagInput.value = cachedSettings.startTag;
endTagInput.value = cachedSettings.endTag;
promptPrefixInput.value = cachedSettings.promptPrefix;
maxWidthInput.value = cachedSettings.maxWidth;
cacheLimitInput.value = cachedSettings.cacheLimit;
apiKeyInput.value = cachedSettings.apiKey;
defaultModelSelect.value = cachedSettings.defaultModel;
document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
await updateCacheStatusDisplay();
}
async function saveSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect) {
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;
cachedSettings.apiKey = apiKeyInput.value.trim();
cachedSettings.defaultModel = defaultModelSelect.value;
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);
await GM_setValue('comfyui_api_key', cachedSettings.apiKey);
await GM_setValue('comfyui_default_model', cachedSettings.defaultModel);
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 ---
function escapeRegex(string) { return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); }
function generateClientId() { return 'client-' + 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} 条旧记录。`);
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();
}
async function cacheImageInBackground(generationId, imageUrl) {
try {
const imageBase64Data = await fetchImageAsBase64(imageUrl);
if (imageBase64Data) {
await saveImageRecord(generationId, imageBase64Data);
}
} catch (e) {
console.error(`后台缓存图片失败 for ${generationId}:`, e);
}
}
// --- Chat Message Processing and Image Generation ---
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, endTag } = cachedSettings;
if (!startTag || !endTag) return;
const regex = new RegExp(
escapeRegex(startTag) +
'(?:\\[model=([\\w.-]+)\\])?' +
'([\\s\\S]*?)' +
escapeRegex(endTag),
'g'
);
const currentHtml = mesText.innerHTML;
if (regex.test(currentHtml) && !mesText.querySelector('.comfy-button-group')) {
mesText.innerHTML = currentHtml.replace(regex, (match, model, prompt) => {
const cleanPrompt = prompt.trim();
const encodedPrompt = cleanPrompt.replace(/"/g, '"');
const modelName = model ? model.trim() : '';
const generationId = simpleHash(modelName + cleanPrompt);
return `<span class="comfy-button-group" data-generation-id="${generationId}"><button class="comfy-button comfy-chat-generate-button" data-prompt="${encodedPrompt}" data-model="${modelName}">开始生成</button></span>`;
});
}
const buttonGroups = mesText.querySelectorAll('.comfy-button-group');
buttonGroups.forEach(group => {
if (group.dataset.listenerAttached) return;
const generationId = group.dataset.generationId;
if (savedImagesCache[generationId]) {
displayImage(group, savedImagesCache[generationId]);
const generateButton = group.querySelector('.comfy-chat-generate-button');
if(generateButton) setupGeneratedState(generateButton, generationId);
}
group.dataset.listenerAttached = 'true';
});
}
function setupGeneratedState(generateButton, generationId) {
generateButton.textContent = '重新生成';
generateButton.disabled = false;
generateButton.classList.remove('testing', 'success', 'error');
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?.classList.contains('comfy-image-container')) {
imageContainer.remove();
}
deleteButton.remove();
generateButton.textContent = '开始生成';
});
generateButton.insertAdjacentElement('afterend', deleteButton);
}
deleteButton.style.display = 'inline-flex';
}
async function onGenerateButtonClickLogic(button) {
const group = button.closest('.comfy-button-group');
let prompt = button.dataset.prompt;
const generationId = group.dataset.generationId;
let model = button.dataset.model || '';
if (!model) {
model = await GM_getValue('comfyui_default_model', '');
}
if (button.disabled) return;
const apiKey = await GM_getValue('comfyui_api_key', '');
if (!apiKey) {
if (typeof toastr !== 'undefined') toastr.error('请先在设置面板中配置 API 密钥！');
document.getElementById(PANEL_ID).style.display = 'flex';
return;
}
if (Date.now() < globalCooldownEndTime) {
const remainingTime = Math.ceil((globalCooldownEndTime - Date.now()) / 1000);
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;
if (oldImageContainer?.classList.contains('comfy-image-container')) {
oldImageContainer.style.opacity = '0.5';
}
try {
connectWebSocket();
const { comfyuiUrl, promptPrefix } = cachedSettings;
if (!comfyuiUrl) throw new Error('调度器 URL 未配置。');
if (promptPrefix) prompt = promptPrefix + ' ' + prompt;
toastr?.info('正在向调度器发送请求...');
const promptResponse = await sendPromptRequestToScheduler(comfyuiUrl, {
client_id: generateClientId(),
positive_prompt: prompt,
api_key: apiKey,
model: model
});
const promptId = promptResponse.prompt_id;
if (!promptId) throw new Error('调度器未返回有效的任务 ID。');
if (socket) {
socket.emit('subscribe_to_prompt', { prompt_id: promptId });
}
activePrompts[promptId] = { button, generationId };
button.textContent = '生成中...';
if(promptResponse.assigned_instance_name) {
toastr?.success(`任务已分配到: ${promptResponse.assigned_instance_name} (队列: ${promptResponse.assigned_instance_queue_size})`);
}
} catch (e) {
console.error('ComfyUI 脚本请求错误:', e);
let displayMessage = e.message || '图片生成失败，请检查服务。';
toastr?.error(displayMessage);
button.textContent = '请求失败';
button.classList.add('error');
setTimeout(() => {
button.classList.remove('testing', 'error');
button.textContent = group.querySelector('.comfy-delete-button') ? '重新生成' : '开始生成';
button.disabled = false;
if(oldImageContainer) oldImageContainer.style.opacity = '1';
if(deleteButton) deleteButton.style.display = 'inline-flex';
}, 3000);
}
}
// --- API Request Functions ---
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: 15000,
onload: (res) => {
if (res.status === 202 || res.status === 200) {
resolve(JSON.parse(res.responseText));
} else {
let errorMsg = `调度器 API 错误: ${res.status}`;
try {
const errorJson = JSON.parse(res.responseText);
if (errorJson.error) errorMsg = errorJson.error;
} catch (e) { /* Do nothing */ }
reject(new Error(errorMsg));
}
},
onerror: (e) => reject(new Error('无法连接到调度器 API。')),
ontimeout: () => reject(new Error('连接调度器 API 超时。')),
});
});
}
// 显示图片，现在可以接受URL或Base64数据
async function displayImage(anchorElement, imageData) {
const group = anchorElement.closest('.comfy-button-group') || anchorElement;
let container = group.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);
group.insertAdjacentElement('afterend', container);
}
container.style.opacity = '1';
const imgElement = container.querySelector('img');
imgElement.src = imageData;
imgElement.style.maxWidth = (cachedSettings.maxWidth || 600) + 'px';
}
// --- 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, {});
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', () => {
loadSettingsFromStorageAndApplyToCache().then(() => {
if (cachedSettings.comfyuiUrl) {
connectWebSocket(); // 页面加载后立即尝试连接
}
});
observeChat();
const body = document.querySelector('body');
if (body) {
optionsObserver.observe(body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style'] });
}
});
})();