移除轮询，使用WebSocket实时接收生成结果，并增加API密钥认证、模型路由和默认模型选择功能。优化版包含安全性增强、智能重连、性能优化等改进。
// ==UserScript==
// @name 公益酒馆ComfyUI插图脚本 (WebSocket实时版 - 优化版)
// @namespace http://tampermonkey.net/
// @version 32.0 // 版本号递增，新增安全性增强、错误处理改进、性能优化等
// @license GPL
// @description 移除轮询，使用WebSocket实时接收生成结果，并增加API密钥认证、模型路由和默认模型选择功能。优化版包含安全性增强、智能重连、性能优化等改进。
// @author feng zheng (升级 by Gemini, 优化 by Claude)
// @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;
const CONFIG_VERSION = '2.0';
const ENCRYPTION_KEY = 42; // 简单的XOR密钥
// --- Security and Utility Functions ---
function encryptApiKey(key) {
if (!key) return '';
try {
return btoa(key.split('').map(c => String.fromCharCode(c.charCodeAt(0) ^ ENCRYPTION_KEY)).join(''));
} catch (e) {
console.error('API密钥加密失败:', e);
return key;
}
}
function decryptApiKey(encryptedKey) {
if (!encryptedKey) return '';
try {
return atob(encryptedKey).split('').map(c => String.fromCharCode(c.charCodeAt(0) ^ ENCRYPTION_KEY)).join('');
} catch (e) {
console.error('API密钥解密失败:', e);
return encryptedKey;
}
}
function sanitizePrompt(prompt) {
if (!prompt || typeof prompt !== 'string') return '';
// 创建临时div元素进行HTML转义
const div = document.createElement('div');
div.textContent = prompt;
return div.innerHTML
.replace(/[<>]/g, '') // 移除尖括号
.replace(/javascript:/gi, '') // 移除javascript:协议
.replace(/on\w+\s*=/gi, '') // 移除事件处理器
.trim();
}
function validateUrl(url) {
if (!url || typeof url !== 'string') return false;
try {
const urlObj = new URL(url);
// 强制要求HTTPS（除了localhost）
if (urlObj.hostname !== 'localhost' && urlObj.hostname !== '127.0.0.1' && urlObj.protocol !== 'https:') {
return false;
}
return ['http:', 'https:'].includes(urlObj.protocol);
} catch (e) {
return false;
}
}
function validateConfig(config) {
const errors = [];
if (!config.comfyuiUrl || typeof config.comfyuiUrl !== 'string') {
errors.push('调度器URL无效');
} else if (!validateUrl(config.comfyuiUrl)) {
errors.push('调度器URL格式错误或不安全');
}
if (config.maxWidth && (typeof config.maxWidth !== 'number' || config.maxWidth < 100 || config.maxWidth > 2000)) {
errors.push('图片最大宽度必须在100-2000像素之间');
}
if (config.cacheLimit && (typeof config.cacheLimit !== 'number' || config.cacheLimit < 1 || config.cacheLimit > 100)) {
errors.push('缓存限制必须在1-100之间');
}
return errors;
}
// --- Error Handling Classes ---
class ComfyUIError extends Error {
constructor(message, type = 'UNKNOWN', details = {}) {
super(message);
this.name = 'ComfyUIError';
this.type = type;
this.details = details;
this.timestamp = Date.now();
}
}
class ErrorHandler {
static handle(error, context = '') {
const errorLog = {
message: error.message || '未知错误',
type: error.type || 'UNKNOWN',
context,
timestamp: error.timestamp || Date.now(),
stack: error.stack
};
console.error('[ComfyUI Error]', errorLog);
const userMessage = this.getUserFriendlyMessage(error.type, error.message);
if (typeof toastr !== 'undefined') {
toastr.error(userMessage);
}
return errorLog;
}
static getUserFriendlyMessage(type, originalMessage) {
const messages = {
'NETWORK': '网络连接失败，请检查网络连接和服务器状态',
'AUTH': '身份验证失败，请检查API密钥是否正确',
'GENERATION': '图片生成失败，请稍后重试',
'CONFIG': '配置错误，请检查设置',
'VALIDATION': '输入验证失败，请检查输入内容',
'CACHE': '缓存操作失败',
'WEBSOCKET': 'WebSocket连接失败，将尝试重新连接'
};
return messages[type] || originalMessage || '操作失败，请重试';
}
}
// --- Smart Reconnection System ---
class SmartReconnector {
constructor(getWsUrl, onConnect, onDisconnect) {
this.getWsUrl = getWsUrl;
this.onConnect = onConnect;
this.onDisconnect = onDisconnect;
this.reconnectAttempts = 0;
this.maxAttempts = 10;
this.baseDelay = 1000;
this.maxDelay = 30000;
this.isOnline = navigator.onLine;
this.setupNetworkMonitoring();
}
setupNetworkMonitoring() {
window.addEventListener('online', () => {
this.isOnline = true;
if (typeof toastr !== 'undefined') {
toastr.success('网络连接已恢复，正在重新连接...');
}
this.reconnect();
});
window.addEventListener('offline', () => {
this.isOnline = false;
if (typeof toastr !== 'undefined') {
toastr.warning('网络连接已断开，将在恢复后自动重连');
}
this.onDisconnect();
});
}
async reconnect() {
if (!this.isOnline) {
console.log('网络离线，暂停重连尝试');
return false;
}
if (this.reconnectAttempts >= this.maxAttempts) {
this.showPermanentDisconnectionNotice();
return false;
}
const delay = Math.min(
this.baseDelay * Math.pow(2, this.reconnectAttempts),
this.maxDelay
);
console.log(`尝试重连 (${this.reconnectAttempts + 1}/${this.maxAttempts})，延迟 ${delay}ms`);
await this.wait(delay);
try {
await this.attemptConnection();
this.reconnectAttempts = 0;
if (typeof toastr !== 'undefined') {
toastr.success('WebSocket连接已恢复！');
}
return true;
} catch (error) {
this.reconnectAttempts++;
console.warn(`重连失败 (${this.reconnectAttempts}/${this.maxAttempts}):`, error.message);
return this.reconnect();
}
}
async attemptConnection() {
return new Promise((resolve, reject) => {
try {
const wsUrl = this.getWsUrl();
if (!wsUrl) {
throw new ComfyUIError('WebSocket URL未配置', 'CONFIG');
}
const testSocket = io(wsUrl, {
timeout: 5000,
reconnection: false
});
const connectTimeout = setTimeout(() => {
testSocket.disconnect();
reject(new ComfyUIError('连接超时', 'WEBSOCKET'));
}, 5000);
testSocket.on('connect', () => {
clearTimeout(connectTimeout);
testSocket.disconnect();
this.onConnect();
resolve();
});
testSocket.on('connect_error', (error) => {
clearTimeout(connectTimeout);
reject(new ComfyUIError('连接失败: ' + error.message, 'WEBSOCKET'));
});
} catch (error) {
reject(new ComfyUIError('连接尝试失败: ' + error.message, 'WEBSOCKET'));
}
});
}
wait(ms) {
return new Promise(resolve => setTimeout(resolve, ms));
}
showPermanentDisconnectionNotice() {
if (typeof toastr !== 'undefined') {
toastr.error('无法连接到服务器，请检查网络和服务器状态');
}
// 在UI中显示离线提示
this.showOfflineNotice();
}
showOfflineNotice() {
// 移除现有的离线提示
const existingNotice = document.querySelector('.comfy-offline-notice');
if (existingNotice) {
existingNotice.remove();
}
const notice = document.createElement('div');
notice.className = 'comfy-offline-notice';
notice.innerHTML = `
<i class="fa fa-wifi-slash"></i>
<span>连接断开，仅显示缓存内容</span>
<button class="retry-connection">重试连接</button>
`;
notice.querySelector('.retry-connection').addEventListener('click', () => {
notice.remove();
this.reconnectAttempts = 0;
this.reconnect();
});
document.body.appendChild(notice);
}
reset() {
this.reconnectAttempts = 0;
}
}
// --- Performance Monitor ---
class PerformanceMonitor {
constructor() {
this.metrics = {};
this.memoryCheckInterval = null;
}
startTimer(operation) {
this.metrics[operation] = performance.now();
}
endTimer(operation) {
if (this.metrics[operation]) {
const duration = performance.now() - this.metrics[operation];
console.log(`[Performance] ${operation}: ${duration.toFixed(2)}ms`);
delete this.metrics[operation];
return duration;
}
return 0;
}
trackMemoryUsage() {
if (performance.memory) {
const usage = {
used: Math.round(performance.memory.usedJSHeapSize / 1024 / 1024),
total: Math.round(performance.memory.totalJSHeapSize / 1024 / 1024),
limit: Math.round(performance.memory.jsHeapSizeLimit / 1024 / 1024)
};
console.log(`[Memory] Used: ${usage.used}MB, Total: ${usage.total}MB, Limit: ${usage.limit}MB`);
// 内存使用超过80%时发出警告
if (usage.used / usage.limit > 0.8) {
console.warn('[Memory Warning] 内存使用过高，建议清理缓存');
if (typeof toastr !== 'undefined') {
toastr.warning('内存使用过高，建议清理图片缓存');
}
}
return usage;
}
return null;
}
startMemoryMonitoring(interval = 30000) {
this.stopMemoryMonitoring();
this.memoryCheckInterval = setInterval(() => {
this.trackMemoryUsage();
}, interval);
}
stopMemoryMonitoring() {
if (this.memoryCheckInterval) {
clearInterval(this.memoryCheckInterval);
this.memoryCheckInterval = null;
}
}
}
// --- Global State Variables ---
let globalCooldownEndTime = 0;
let socket = null;
let activePrompts = {}; // 存储 prompt_id -> { button, generationId } 的映射
let reconnector = null;
let performanceMonitor = new PerformanceMonitor();
let cachedDOMElements = {}; // DOM元素缓存
let debugMode = false; // 调试模式开关
// --- Cached User Settings ---
let cachedSettings = {
comfyuiUrl: '',
startTag: 'image###',
endTag: '###',
promptPrefix: '',
maxWidth: 600,
cacheLimit: 20,
apiKey: '', // 将存储加密后的密钥
defaultModel: ''
};
// --- Inject Custom CSS Styles ---
GM_addStyle(`
/* 新增：离线提示样式 */
.comfy-offline-notice {
position: fixed;
top: 20px;
right: 20px;
background: rgba(220, 53, 69, 0.9);
color: white;
padding: 10px 15px;
border-radius: 5px;
display: flex;
align-items: center;
gap: 10px;
z-index: 10000;
font-size: 14px;
box-shadow: 0 2px 10px rgba(0,0,0,0.3);
}
.comfy-offline-notice .retry-connection {
background: rgba(255,255,255,0.2);
border: 1px solid rgba(255,255,255,0.3);
color: white;
padding: 5px 10px;
border-radius: 3px;
cursor: pointer;
font-size: 12px;
}
.comfy-offline-notice .retry-connection:hover {
background: rgba(255,255,255,0.3);
}
/* 新增：缓存状态显示样式 */
#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;
}
/* 新增：配置验证错误提示 */
.comfy-config-error {
background-color: rgba(220, 53, 69, 0.1);
border: 1px solid #dc3545;
color: #dc3545;
padding: 8px;
border-radius: 4px;
margin: 10px 0;
font-size: 0.9em;
}
/* 控制面板主容器样式 */
#${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;
}
`);
// --- Configuration Migration ---
async function migrateConfig() {
try {
const currentVersion = await GM_getValue('config_version', '1.0');
if (currentVersion !== CONFIG_VERSION) {
console.log(`配置迁移: ${currentVersion} -> ${CONFIG_VERSION}`);
// 迁移旧版本的API密钥
const oldApiKey = await GM_getValue('comfyui_api_key', '');
if (oldApiKey && !oldApiKey.includes('=')) { // 检查是否已加密
const encryptedKey = encryptApiKey(oldApiKey);
await GM_setValue('comfyui_api_key', encryptedKey);
console.log('API密钥已加密存储');
}
await GM_setValue('config_version', CONFIG_VERSION);
console.log('配置迁移完成');
}
} catch (error) {
console.error('配置迁移失败:', error);
}
}
// --- Cache Integrity and Management ---
async function validateCacheIntegrity() {
try {
performanceMonitor.startTimer('validateCache');
const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
const validRecords = {};
let removedCount = 0;
for (const [id, data] of Object.entries(records)) {
try {
// 验证Base64数据完整性
if (typeof data === 'string' &&
data.startsWith('data:image/') &&
data.includes('base64,') &&
data.length > 100) { // 基本长度检查
validRecords[id] = data;
} else {
console.warn(`缓存记录 ${id} 数据格式无效，已清理`);
removedCount++;
}
} catch (error) {
console.error(`缓存记录 ${id} 验证失败:`, error);
removedCount++;
}
}
if (removedCount > 0) {
await GM_setValue(STORAGE_KEY_IMAGES, validRecords);
if (typeof toastr !== 'undefined') {
toastr.info(`已清理 ${removedCount} 条无效缓存记录`);
}
}
performanceMonitor.endTimer('validateCache');
return Object.keys(validRecords).length;
} catch (error) {
ErrorHandler.handle(new ComfyUIError('缓存验证失败: ' + error.message, 'CACHE'), 'validateCacheIntegrity');
return 0;
}
}
// --- Image Compression and Optimization ---
function compressImage(canvas, quality = 0.8, maxWidth = 1024) {
return new Promise((resolve) => {
try {
// 如果图片太大，先缩放
if (canvas.width > maxWidth) {
const scale = maxWidth / canvas.width;
const scaledCanvas = document.createElement('canvas');
const ctx = scaledCanvas.getContext('2d');
scaledCanvas.width = maxWidth;
scaledCanvas.height = canvas.height * scale;
ctx.drawImage(canvas, 0, 0, scaledCanvas.width, scaledCanvas.height);
canvas = scaledCanvas;
}
// 压缩为JPEG格式
canvas.toBlob((blob) => {
if (blob) {
const reader = new FileReader();
reader.onload = () => resolve(reader.result);
reader.readAsDataURL(blob);
} else {
resolve(null);
}
}, 'image/jpeg', quality);
} catch (error) {
console.error('图片压缩失败:', error);
resolve(null);
}
});
}
async function fetchImageAsBase64(imageUrl) {
return new Promise((resolve, reject) => {
performanceMonitor.startTimer('fetchImage');
GM_xmlhttpRequest({
method: 'GET',
url: imageUrl,
responseType: 'blob',
timeout: 30000,
onload: async (response) => {
try {
if (response.status === 200) {
const blob = response.response;
// 检查文件大小
if (blob.size > 10 * 1024 * 1024) { // 10MB限制
throw new ComfyUIError('图片文件过大', 'VALIDATION');
}
// 尝试压缩大图片
if (blob.size > 2 * 1024 * 1024) { // 2MB以上尝试压缩
const img = new Image();
const canvas = document.createElement('canvas');
const ctx = canvas.getContext('2d');
img.onload = async () => {
canvas.width = img.width;
canvas.height = img.height;
ctx.drawImage(img, 0, 0);
const compressedData = await compressImage(canvas, 0.8, 1024);
if (compressedData) {
console.log(`图片已压缩: ${(blob.size / 1024).toFixed(1)}KB -> ${(compressedData.length * 0.75 / 1024).toFixed(1)}KB`);
resolve(compressedData);
} else {
// 压缩失败，使用原图
const reader = new FileReader();
reader.onloadend = () => resolve(reader.result);
reader.readAsDataURL(blob);
}
};
img.onerror = () => {
// 图片解析失败，使用原始数据
const reader = new FileReader();
reader.onloadend = () => resolve(reader.result);
reader.readAsDataURL(blob);
};
const reader = new FileReader();
reader.onload = () => img.src = reader.result;
reader.readAsDataURL(blob);
} else {
// 小文件直接转换
const reader = new FileReader();
reader.onloadend = () => resolve(reader.result);
reader.onerror = (err) => reject(new ComfyUIError('FileReader error: ' + err, 'CACHE'));
reader.readAsDataURL(blob);
}
} else {
reject(new ComfyUIError(`获取图片失败，状态: ${response.status}`, 'NETWORK'));
}
} catch (error) {
reject(error);
} finally {
performanceMonitor.endTimer('fetchImage');
}
},
onerror: (err) => {
performanceMonitor.endTimer('fetchImage');
reject(new ComfyUIError('网络错误: ' + err, 'NETWORK'));
},
ontimeout: () => {
performanceMonitor.endTimer('fetchImage');
reject(new ComfyUIError('下载图片超时', 'NETWORK'));
}
});
});
}
// --- WebSocket & State Management ---
function connectWebSocket() {
try {
if (socket && socket.connected) return;
const schedulerUrl = new URL(cachedSettings.comfyuiUrl);
const wsUrl = `${schedulerUrl.protocol}//${schedulerUrl.host}`;
if (typeof io === 'undefined') {
throw new ComfyUIError('Socket.IO 客户端库未加载！', 'CONFIG');
}
socket = io(wsUrl, {
reconnectionAttempts: 5,
timeout: 20000,
transports: ['websocket', 'polling'] // 添加备用传输方式
});
socket.on('connect', () => {
console.log('成功连接到调度器 WebSocket！');
if (typeof toastr !== 'undefined') toastr.success('已建立实时连接！');
// 重置重连器状态
if (reconnector) {
reconnector.reset();
}
// 移除离线提示
const offlineNotice = document.querySelector('.comfy-offline-notice');
if (offlineNotice) {
offlineNotice.remove();
}
});
socket.on('disconnect', (reason) => {
console.log('与调度器 WebSocket 断开连接:', reason);
// 只在非主动断开时尝试重连
if (reason !== 'io client disconnect' && reason !== 'io server disconnect') {
if (reconnector) {
reconnector.reconnect();
}
}
});
socket.on('connect_error', (error) => {
console.error('WebSocket连接错误:', error);
ErrorHandler.handle(new ComfyUIError('WebSocket连接失败: ' + error.message, 'WEBSOCKET'), 'connectWebSocket');
});
socket.on('generation_complete', (data) => {
try {
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];
} catch (error) {
ErrorHandler.handle(new ComfyUIError('处理生成完成事件失败: ' + error.message, 'WEBSOCKET'), 'generation_complete');
}
});
} catch (error) {
ErrorHandler.handle(error, 'connectWebSocket');
}
}
// --- 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">
<div id="comfyui-config-errors" class="comfy-config-error" style="display: none;"></div>
<label for="comfyui-url">调度器 URL (推荐使用HTTPS)</label>
<div class="comfy-url-container">
<input id="comfyui-url" type="text" placeholder="例如: https://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="password" 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, 100-2000):</label><input id="comfyui-max-width" type="number" placeholder="例如: 600" min="100" max="2000"></div>
<div><label for="comfyui-cache-limit">最大缓存数量 (1-100):</label><input id="comfyui-cache-limit" type="number" placeholder="例如: 20" min="1" max="100"></div>
<div id="comfyui-cache-status">当前缓存: ...</div>
<div style="margin-top: 15px;">
<label>
<input type="checkbox" id="comfyui-debug-mode" style="width: auto; margin-right: 8px;">
启用调试模式（在控制台查看详细日志）
</label>
</div>
<button id="comfyui-force-rescan" class="comfy-button" style="margin-top: 10px; width: 100%;">强制重新扫描所有消息</button>
<button id="comfyui-clear-cache" class="comfy-button error" style="margin-top: 10px; width: 100%;">删除所有图片缓存</button>
<button id="comfyui-validate-cache" class="comfy-button" style="margin-top: 10px; width: 100%;">验证缓存完整性</button>
</div>
</div>
`;
document.body.insertAdjacentHTML('beforeend', panelHTML);
initPanelLogic();
}
async function updateCacheStatusDisplay() {
try {
const display = getCachedDOMElement('comfyui-cache-status') || document.getElementById('comfyui-cache-status');
if (!display) return;
const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
const count = Object.keys(records).length;
const totalSize = Object.values(records).reduce((total, data) => {
return total + (typeof data === 'string' ? data.length : 0);
}, 0);
const sizeMB = (totalSize * 0.75 / 1024 / 1024).toFixed(1); // Base64编码大约增大33%
display.textContent = `当前缓存: ${count} / ${cachedSettings.cacheLimit} 张 (约 ${sizeMB}MB)`;
setCachedDOMElement('comfyui-cache-status', display);
} catch (error) {
ErrorHandler.handle(new ComfyUIError('更新缓存状态失败: ' + error.message, 'CACHE'), 'updateCacheStatusDisplay');
}
}
// DOM缓存管理
function getCachedDOMElement(id) {
return cachedDOMElements[id];
}
function setCachedDOMElement(id, element) {
cachedDOMElements[id] = element;
}
function clearDOMCache() {
cachedDOMElements = {};
}
function showConfigErrors(errors) {
const errorContainer = document.getElementById('comfyui-config-errors');
if (errorContainer) {
if (errors.length > 0) {
errorContainer.innerHTML = errors.map(error => `• ${error}`).join('<br>');
errorContainer.style.display = 'block';
} else {
errorContainer.style.display = 'none';
}
}
}
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 validateCacheButton = document.getElementById('comfyui-validate-cache');
const forceRescanButton = document.getElementById('comfyui-force-rescan');
const debugModeCheckbox = document.getElementById('comfyui-debug-mode');
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');
// 缓存常用DOM元素
setCachedDOMElement('panel', panel);
setCachedDOMElement('testButton', testButton);
closeButton.addEventListener('click', () => {
panel.style.display = 'none';
});
testButton.addEventListener('click', () => {
try {
let url = urlInput.value.trim();
if (!url) {
throw new ComfyUIError('请输入调度器的URL', 'VALIDATION');
}
if (!url.startsWith('http')) {
url = 'https://' + url; // 默认使用HTTPS
}
if (url.endsWith('/')) {
url = url.slice(0, -1);
}
// 验证URL
if (!validateUrl(url)) {
throw new ComfyUIError('URL格式无效或不安全，请使用HTTPS', 'VALIDATION');
}
urlInput.value = url;
const testUrl = url + '/system_stats';
if (typeof toastr !== 'undefined') toastr.info('正在尝试连接服务...');
testButton.className = 'comfy-button testing';
testButton.disabled = true;
performanceMonitor.startTimer('testConnection');
GM_xmlhttpRequest({
method: "GET",
url: testUrl,
timeout: 10000,
onload: (res) => {
performanceMonitor.endTimer('testConnection');
testButton.disabled = false;
testButton.className = res.status === 200 ? 'comfy-button success' : 'comfy-button error';
if (res.status === 200) {
if (typeof toastr !== 'undefined') toastr.success('连接成功！');
} else {
if (typeof toastr !== 'undefined') toastr.error(`连接失败！状态: ${res.status}`);
}
},
onerror: (error) => {
performanceMonitor.endTimer('testConnection');
testButton.disabled = false;
testButton.className = 'comfy-button error';
ErrorHandler.handle(new ComfyUIError('连接失败', 'NETWORK'), 'testConnection');
},
ontimeout: () => {
performanceMonitor.endTimer('testConnection');
testButton.disabled = false;
testButton.className = 'comfy-button error';
ErrorHandler.handle(new ComfyUIError('连接超时', 'NETWORK'), 'testConnection');
}
});
} catch (error) {
ErrorHandler.handle(error, 'testConnection');
testButton.disabled = false;
testButton.className = 'comfy-button error';
}
});
clearCacheButton.addEventListener('click', async () => {
if (confirm('您确定要删除所有已生成的图片缓存吗？')) {
try {
performanceMonitor.startTimer('clearCache');
await GM_setValue(STORAGE_KEY_IMAGES, {});
await updateCacheStatusDisplay();
performanceMonitor.endTimer('clearCache');
if (typeof toastr !== 'undefined') toastr.success('图片缓存已清空！');
} catch (error) {
ErrorHandler.handle(new ComfyUIError('清空缓存失败: ' + error.message, 'CACHE'), 'clearCache');
}
}
});
validateCacheButton.addEventListener('click', async () => {
try {
validateCacheButton.disabled = true;
validateCacheButton.textContent = '验证中...';
const validCount = await validateCacheIntegrity();
await updateCacheStatusDisplay();
validateCacheButton.disabled = false;
validateCacheButton.textContent = '验证缓存完整性';
if (typeof toastr !== 'undefined') {
toastr.success(`缓存验证完成，有效记录: ${validCount} 条`);
}
} catch (error) {
validateCacheButton.disabled = false;
validateCacheButton.textContent = '验证缓存完整性';
ErrorHandler.handle(error, 'validateCache');
}
});
// 调试模式开关
debugModeCheckbox.addEventListener('change', async () => {
debugMode = debugModeCheckbox.checked;
await GM_setValue('comfyui_debug_mode', debugMode);
if (debugMode) {
console.log('[ComfyUI Debug] 调试模式已启用');
if (typeof toastr !== 'undefined') {
toastr.info('调试模式已启用，请查看浏览器控制台');
}
} else {
console.log('[ComfyUI Debug] 调试模式已禁用');
}
});
// 强制重新扫描所有消息
forceRescanButton.addEventListener('click', async () => {
try {
forceRescanButton.disabled = true;
forceRescanButton.textContent = '扫描中...';
const chatElement = document.getElementById('chat');
if (chatElement) {
const allMessages = chatElement.querySelectorAll('.mes');
const savedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});
console.log(`[ComfyUI] 开始强制重新扫描 ${allMessages.length} 条消息`);
// 清除已有的监听器标记
allMessages.forEach(node => {
const mesTextElement = node.querySelector('.mes_text');
if (mesTextElement) {
mesTextElement.dataset.listenersAttached = '';
// 移除现有的按钮组以避免重复
const existingButtons = mesTextElement.querySelectorAll('.comfy-button-group');
existingButtons.forEach(btn => btn.remove());
}
});
// 重新处理所有消息
for (const node of allMessages) {
try {
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';
}
await processMessageForComfyButton(node, savedImages);
} catch (error) {
console.error('处理消息节点失败:', error);
}
}
console.log('[ComfyUI] 强制重新扫描完成');
if (typeof toastr !== 'undefined') {
toastr.success(`已重新扫描 ${allMessages.length} 条消息`);
}
} else {
if (typeof toastr !== 'undefined') {
toastr.error('未找到聊天区域，无法执行扫描');
}
}
forceRescanButton.disabled = false;
forceRescanButton.textContent = '强制重新扫描所有消息';
} catch (error) {
forceRescanButton.disabled = false;
forceRescanButton.textContent = '强制重新扫描所有消息';
ErrorHandler.handle(new ComfyUIError('强制重新扫描失败: ' + error.message, 'UI'), 'forceRescan');
}
});
loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect, debugModeCheckbox).then(() => {
applyCurrentMaxWidthToAllImages();
});
[urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect, debugModeCheckbox].forEach(input => {
const eventType = input.type === 'checkbox' ? 'change' : (input.tagName.toLowerCase() === 'select' ? 'change' : 'input');
input.addEventListener(eventType, async () => {
try {
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();
setTimeout(connectWebSocket, 500); // 延迟重连
}
} catch (error) {
ErrorHandler.handle(error, 'saveSettings');
}
});
});
}
async function loadSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect, debugModeCheckbox) {
try {
performanceMonitor.startTimer('loadSettings');
cachedSettings.comfyuiUrl = await GM_getValue('comfyui_url', 'https://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);
// 解密API密钥
const encryptedApiKey = await GM_getValue('comfyui_api_key', '');
cachedSettings.apiKey = decryptApiKey(encryptedApiKey);
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;
// 加载调试模式设置
debugMode = await GM_getValue('comfyui_debug_mode', false);
debugModeCheckbox.checked = debugMode;
document.documentElement.style.setProperty('--comfy-image-max-width', (cachedSettings.maxWidth || 600) + 'px');
await updateCacheStatusDisplay();
// 验证配置
const configErrors = validateConfig(cachedSettings);
showConfigErrors(configErrors);
performanceMonitor.endTimer('loadSettings');
} catch (error) {
ErrorHandler.handle(new ComfyUIError('加载设置失败: ' + error.message, 'CONFIG'), 'loadSettings');
}
}
async function saveSettings(urlInput, startTagInput, endTagInput, promptPrefixInput, maxWidthInput, cacheLimitInput, apiKeyInput, defaultModelSelect) {
try {
performanceMonitor.startTimer('saveSettings');
const newSettings = {
comfyuiUrl: urlInput.value.trim(),
startTag: startTagInput.value,
endTag: endTagInput.value,
promptPrefix: promptPrefixInput.value.trim(),
maxWidth: parseInt(maxWidthInput.value) || 600,
cacheLimit: parseInt(cacheLimitInput.value) || 20,
apiKey: apiKeyInput.value.trim(),
defaultModel: defaultModelSelect.value
};
// 验证新配置
const configErrors = validateConfig(newSettings);
showConfigErrors(configErrors);
// 如果有严重错误，不保存配置
if (configErrors.length > 0 && configErrors.some(error => error.includes('URL'))) {
return;
}
// 更新缓存设置
Object.assign(cachedSettings, newSettings);
// 加密并保存API密钥
const encryptedApiKey = encryptApiKey(cachedSettings.apiKey);
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', encryptedApiKey);
await GM_setValue('comfyui_default_model', cachedSettings.defaultModel);
document.documentElement.style.setProperty('--comfy-image-max-width', cachedSettings.maxWidth + 'px');
await updateCacheStatusDisplay();
performanceMonitor.endTimer('saveSettings');
} catch (error) {
ErrorHandler.handle(new ComfyUIError('保存设置失败: ' + error.message, 'CONFIG'), 'saveSettings');
}
}
async function applyCurrentMaxWidthToAllImages() {
try {
const images = document.querySelectorAll('.comfy-image-container img');
const maxWidthPx = (cachedSettings.maxWidth || 600) + 'px';
images.forEach(img => {
img.style.maxWidth = maxWidthPx;
});
} catch (error) {
console.error('应用图片宽度设置失败:', error);
}
}
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) + '-' + Date.now();
}
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);
}
async function saveImageRecord(generationId, imageBase64Data) {
try {
performanceMonitor.startTimer('saveImageRecord');
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();
performanceMonitor.endTimer('saveImageRecord');
} catch (error) {
ErrorHandler.handle(new ComfyUIError('保存图片记录失败: ' + error.message, 'CACHE'), 'saveImageRecord');
}
}
async function deleteImageRecord(generationId) {
try {
const records = await GM_getValue(STORAGE_KEY_IMAGES, {});
delete records[generationId];
await GM_setValue(STORAGE_KEY_IMAGES, records);
await updateCacheStatusDisplay();
} catch (error) {
ErrorHandler.handle(new ComfyUIError('删除图片记录失败: ' + error.message, 'CACHE'), 'deleteImageRecord');
}
}
async function cacheImageInBackground(generationId, imageUrl) {
try {
const imageBase64Data = await fetchImageAsBase64(imageUrl);
if (imageBase64Data) {
await saveImageRecord(generationId, imageBase64Data);
}
} catch (error) {
console.error(`后台缓存图片失败 for ${generationId}:`, error);
// 缓存失败不影响主流程，只记录错误
}
}
// --- Chat Message Processing and Image Generation ---
function handleComfyButtonClick(event, isTouch = false) {
try {
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);
}
} catch (error) {
ErrorHandler.handle(new ComfyUIError('处理按钮点击失败: ' + error.message, 'UI'), 'handleComfyButtonClick');
}
}
async function processMessageForComfyButton(messageNode, savedImagesCache) {
try {
const mesText = messageNode.querySelector('.mes_text');
if (!mesText) {
if (debugMode) console.log('[ComfyUI Debug] 未找到 .mes_text 元素');
return;
}
const { startTag, endTag } = cachedSettings;
if (!startTag || !endTag) {
if (debugMode) console.log('[ComfyUI Debug] 开始或结束标签未配置:', { startTag, endTag });
return;
}
if (debugMode) {
console.log('[ComfyUI Debug] 开始处理消息, 标签:', { startTag, endTag });
console.log('[ComfyUI Debug] 消息内容:', mesText.innerHTML);
}
const regex = new RegExp(
escapeRegex(startTag) +
'(?:\\[model=([\\w.-]+)\\])?' +
'([\\s\\S]*?)' +
escapeRegex(endTag),
'g'
);
if (debugMode) console.log('[ComfyUI Debug] 使用的正则表达式:', regex);
const currentHtml = mesText.innerHTML;
const matches = currentHtml.match(regex);
if (debugMode) console.log('[ComfyUI Debug] 正则匹配结果:', matches);
// 检查是否存在文本内容中的标签（未被HTML转义的）
const textContent = mesText.textContent || mesText.innerText || '';
if (debugMode) console.log('[ComfyUI Debug] 纯文本内容:', textContent);
const textMatches = textContent.match(regex);
if (debugMode) console.log('[ComfyUI Debug] 纯文本匹配结果:', textMatches);
// 尝试处理HTML内容
let processedHtml = false;
if (regex.test(currentHtml) && !mesText.querySelector('.comfy-button-group')) {
if (debugMode) console.log('[ComfyUI Debug] 在HTML中发现匹配，开始替换');
mesText.innerHTML = currentHtml.replace(regex, (match, model, prompt) => {
if (debugMode) console.log('[ComfyUI Debug] 替换匹配项:', { match, model, prompt });
const cleanPrompt = sanitizePrompt(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>`;
});
processedHtml = true;
}
// 如果HTML中没有找到，尝试处理纯文本内容
if (!processedHtml && textMatches && textMatches.length > 0 && !mesText.querySelector('.comfy-button-group')) {
if (debugMode) console.log('[ComfyUI Debug] HTML中未找到匹配，尝试处理纯文本内容');
// 重置正则表达式状态
regex.lastIndex = 0;
let newHtml = currentHtml;
let match;
while ((match = regex.exec(textContent)) !== null) {
if (debugMode) console.log('[ComfyUI Debug] 处理纯文本匹配:', match);
const fullMatch = match[0];
const model = match[1] || '';
const prompt = match[2] || '';
const cleanPrompt = sanitizePrompt(prompt.trim());
const encodedPrompt = cleanPrompt.replace(/"/g, '"');
const modelName = model.trim();
const generationId = simpleHash(modelName + cleanPrompt);
const buttonHtml = `<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>`;
// 在HTML中查找并替换对应的文本
newHtml = newHtml.replace(escapeRegex(fullMatch), buttonHtml);
}
if (newHtml !== currentHtml) {
mesText.innerHTML = newHtml;
processedHtml = true;
}
}
// 特殊处理：检查是否有被HTML编码的标签
const htmlEncodedStartTag = startTag.replace(/</g, '<').replace(/>/g, '>');
const htmlEncodedEndTag = endTag.replace(/</g, '<').replace(/>/g, '>');
if (htmlEncodedStartTag !== startTag || htmlEncodedEndTag !== endTag) {
if (debugMode) console.log('[ComfyUI Debug] 检查HTML编码的标签:', { htmlEncodedStartTag, htmlEncodedEndTag });
const htmlEncodedRegex = new RegExp(
escapeRegex(htmlEncodedStartTag) +
'(?:\\[model=([\\w.-]+)\\])?' +
'([\\s\\S]*?)' +
escapeRegex(htmlEncodedEndTag),
'g'
);
if (htmlEncodedRegex.test(currentHtml) && !mesText.querySelector('.comfy-button-group')) {
if (debugMode) console.log('[ComfyUI Debug] 发现HTML编码的标签，进行替换');
mesText.innerHTML = currentHtml.replace(htmlEncodedRegex, (match, model, prompt) => {
if (debugMode) console.log('[ComfyUI Debug] 替换HTML编码匹配项:', { match, model, prompt });
const cleanPrompt = sanitizePrompt(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>`;
});
processedHtml = true;
}
}
if (debugMode && !processedHtml) {
console.log('[ComfyUI Debug] 未找到任何匹配的标签');
}
const buttonGroups = mesText.querySelectorAll('.comfy-button-group');
if (debugMode) console.log('[ComfyUI Debug] 找到的按钮组数量:', buttonGroups.length);
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';
});
} catch (error) {
ErrorHandler.handle(new ComfyUIError('处理消息失败: ' + error.message, 'UI'), 'processMessageForComfyButton');
}
}
function setupGeneratedState(generateButton, generationId) {
try {
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 () => {
try {
await deleteImageRecord(generationId);
const imageContainer = group.nextElementSibling;
if (imageContainer?.classList.contains('comfy-image-container')) {
imageContainer.remove();
}
deleteButton.remove();
generateButton.textContent = '开始生成';
} catch (error) {
ErrorHandler.handle(new ComfyUIError('删除图片失败: ' + error.message, 'CACHE'), 'deleteImage');
}
});
generateButton.insertAdjacentElement('afterend', deleteButton);
}
deleteButton.style.display = 'inline-flex';
} catch (error) {
ErrorHandler.handle(new ComfyUIError('设置生成状态失败: ' + error.message, 'UI'), 'setupGeneratedState');
}
}
async function onGenerateButtonClickLogic(button) {
try {
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;
// 获取并解密API密钥
const encryptedApiKey = await GM_getValue('comfyui_api_key', '');
const apiKey = decryptApiKey(encryptedApiKey);
if (!apiKey) {
throw new ComfyUIError('请先在设置面板中配置 API 密钥！', 'AUTH');
}
if (Date.now() < globalCooldownEndTime) {
const remainingTime = Math.ceil((globalCooldownEndTime - Date.now()) / 1000);
if (typeof toastr !== 'undefined') toastr.warning(`请稍候，冷却中 (${remainingTime}s)。`);
return;
}
// 验证和净化提示词
prompt = sanitizePrompt(prompt);
if (!prompt) {
throw new ComfyUIError('提示词内容无效', 'VALIDATION');
}
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';
}
performanceMonitor.startTimer('generateImage');
connectWebSocket();
const { comfyuiUrl, promptPrefix } = cachedSettings;
if (!comfyuiUrl) {
throw new ComfyUIError('调度器 URL 未配置', 'CONFIG');
}
if (!validateUrl(comfyuiUrl)) {
throw new ComfyUIError('调度器 URL 格式无效', 'CONFIG');
}
if (promptPrefix) prompt = promptPrefix + ' ' + prompt;
if (typeof toastr !== 'undefined') 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 ComfyUIError('调度器未返回有效的任务 ID', 'GENERATION');
}
if (socket && socket.connected) {
socket.emit('subscribe_to_prompt', { prompt_id: promptId });
} else {
throw new ComfyUIError('WebSocket连接未建立', 'WEBSOCKET');
}
activePrompts[promptId] = { button, generationId };
button.textContent = '生成中...';
if(promptResponse.assigned_instance_name) {
if (typeof toastr !== 'undefined') {
toastr.success(`任务已分配到: ${promptResponse.assigned_instance_name} (队列: ${promptResponse.assigned_instance_queue_size})`);
}
}
performanceMonitor.endTimer('generateImage');
} catch (error) {
performanceMonitor.endTimer('generateImage');
ErrorHandler.handle(error, 'generateImage');
button.textContent = error.type === 'AUTH' ? '认证失败' : '请求失败';
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);
if (error.type === 'AUTH') {
// 认证失败时打开设置面板
document.getElementById(PANEL_ID).style.display = 'flex';
}
}
}
// --- API Request Functions ---
function sendPromptRequestToScheduler(url, payload) {
return new Promise((resolve, reject) => {
// 验证payload
if (!payload.api_key) {
reject(new ComfyUIError('API密钥缺失', 'AUTH'));
return;
}
if (!payload.positive_prompt || payload.positive_prompt.trim().length === 0) {
reject(new ComfyUIError('提示词不能为空', 'VALIDATION'));
return;
}
GM_xmlhttpRequest({
method: 'POST',
url: `${url}/generate`,
headers: {
'Content-Type': 'application/json',
'User-Agent': 'ComfyUI-TamperMonkey-Script/2.0'
},
data: JSON.stringify(payload),
timeout: 15000,
onload: (res) => {
try {
if (res.status === 202 || res.status === 200) {
const responseData = JSON.parse(res.responseText);
resolve(responseData);
} else if (res.status === 401 || res.status === 403) {
reject(new ComfyUIError('API密钥无效或权限不足', 'AUTH'));
} else {
let errorMsg = `调度器 API 错误: ${res.status}`;
try {
const errorJson = JSON.parse(res.responseText);
if (errorJson.error) errorMsg = errorJson.error;
} catch (e) {
// JSON解析失败，使用默认错误消息
}
reject(new ComfyUIError(errorMsg, 'GENERATION'));
}
} catch (error) {
reject(new ComfyUIError('解析服务器响应失败: ' + error.message, 'NETWORK'));
}
},
onerror: (e) => reject(new ComfyUIError('无法连接到调度器 API', 'NETWORK')),
ontimeout: () => reject(new ComfyUIError('连接调度器 API 超时', 'NETWORK')),
});
});
}
// 显示图片，现在可以接受URL或Base64数据
async function displayImage(anchorElement, imageData) {
try {
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 生成的图片';
img.loading = 'lazy'; // 懒加载优化
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';
} catch (error) {
ErrorHandler.handle(new ComfyUIError('显示图片失败: ' + error.message, 'UI'), 'displayImage');
}
}
// --- Main Execution Logic ---
createComfyUIPanel();
const chatObserver = new MutationObserver(async (mutations) => {
try {
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 => {
try {
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);
} catch (error) {
console.error('处理消息节点失败:', error);
}
});
}
} catch (error) {
ErrorHandler.handle(new ComfyUIError('聊天观察器处理失败: ' + error.message, 'UI'), 'chatObserver');
}
});
async function loadSettingsFromStorageAndApplyToCache() {
try {
await migrateConfig(); // 执行配置迁移
cachedSettings.comfyuiUrl = await GM_getValue('comfyui_url', 'https://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');
} catch (error) {
ErrorHandler.handle(new ComfyUIError('加载初始设置失败: ' + error.message, 'CONFIG'), 'loadSettingsFromStorageAndApplyToCache');
}
}
function observeChat() {
const chatElement = document.getElementById('chat');
if (chatElement) {
loadSettingsFromStorageAndApplyToCache().then(async () => {
try {
const initialSavedImages = await GM_getValue(STORAGE_KEY_IMAGES, {});
chatElement.querySelectorAll('.mes').forEach(node => {
try {
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);
} catch (error) {
console.error('初始化消息节点失败:', error);
}
});
chatObserver.observe(chatElement, { childList: true, subtree: true });
} catch (error) {
ErrorHandler.handle(new ComfyUIError('初始化聊天观察失败: ' + error.message, 'UI'), 'observeChat');
}
});
} else {
setTimeout(observeChat, 500);
}
}
const optionsObserver = new MutationObserver(() => {
try {
const optionsMenu = document.getElementById('options');
if (optionsMenu && optionsMenu.style.display !== 'none') {
addMainButton();
}
} catch (error) {
console.error('选项观察器错误:', error);
}
});
// 初始化重连器
function initializeReconnector() {
reconnector = new SmartReconnector(
() => {
if (cachedSettings.comfyuiUrl) {
const schedulerUrl = new URL(cachedSettings.comfyuiUrl);
return `${schedulerUrl.protocol}//${schedulerUrl.host}`;
}
return null;
},
connectWebSocket,
() => {
if (socket) {
socket.disconnect();
socket = null;
}
}
);
}
// 页面卸载时清理资源
window.addEventListener('beforeunload', () => {
try {
performanceMonitor.stopMemoryMonitoring();
if (socket) {
socket.disconnect();
}
clearDOMCache();
} catch (error) {
console.error('清理资源失败:', error);
}
});
window.addEventListener('load', () => {
try {
loadSettingsFromStorageAndApplyToCache().then(() => {
initializeReconnector();
if (cachedSettings.comfyuiUrl && validateUrl(cachedSettings.comfyuiUrl)) {
connectWebSocket(); // 页面加载后立即尝试连接
}
// 启动性能监控
performanceMonitor.startMemoryMonitoring();
// 定期验证缓存完整性
setInterval(async () => {
try {
await validateCacheIntegrity();
} catch (error) {
console.error('定期缓存验证失败:', error);
}
}, 300000); // 每5分钟验证一次
}).catch(error => {
ErrorHandler.handle(new ComfyUIError('初始化失败: ' + error.message, 'CONFIG'), 'window.load');
});
observeChat();
const body = document.querySelector('body');
if (body) {
optionsObserver.observe(body, {
childList: true,
subtree: true,
attributes: true,
attributeFilter: ['style']
});
}
} catch (error) {
ErrorHandler.handle(new ComfyUIError('页面加载初始化失败: ' + error.message, 'CONFIG'), 'window.load');
}
});
})();